"""W&B metrics, checkpoint artifacts and run continuation."""

from contextlib import contextmanager

from .checkpoint import save_checkpoint
from .provenance import VERSION


@contextmanager
def open_run(config, directory, model, datasets, previous=None):
    """Create a tracking context; disabled mode does not initialize W&B."""
    settings = config.get("wandb", {})
    mode = settings.get("mode", "disabled")
    interval = settings.get("log_every", 10)
    if mode not in ("online", "offline", "disabled"):
        raise ValueError("wandb.mode must be online, offline, or disabled")
    if isinstance(interval, bool) or not isinstance(interval, int) or interval < 1:
        raise ValueError("wandb.log_every must be a positive integer")
    if mode == "disabled":
        yield None
        return

    import wandb
    project = settings.get("project", "paya")
    entity = settings.get("entity")
    continuation = {}
    if previous and mode == previous["mode"] == "online":
        # W&B can return canonical capitalization different from the requested project.
        if project.casefold() != previous["project"].casefold() or entity not in (None, previous["entity"]):
            raise ValueError("Online resume must use the saved W&B project and entity")
        project = previous["project"]
        entity = previous["entity"]
        continuation = {"id": previous["id"], "resume": "must"}
    with wandb.init(project=project, entity=entity, name=settings.get("name") or directory.name,
                     mode=mode, dir=str(directory), save_code=False, **continuation,
                     settings=wandb.Settings(init_timeout=90,finish_timeout=settings.get("finish_timeout",120),
                                             finish_timeout_raises=True),
                     config=None if continuation else {"version": VERSION, "model": model.identity(),
                             "training": {k: v for k, v in config.items()
                                          if k not in ("wandb", "data", "assets", "feature_cache")},
                             "datasets": {name: {"records": len(data.rows), "sha256": data.sha256}
                                          for name, data in datasets.items()},
                             "resumed_from_run_id": previous["id"] if previous else None}) as run:
        run.define_metric("step")
        run.define_metric("epoch")
        run.define_metric("train/*", step_metric="step")
        run.define_metric("selection/*", step_metric="epoch")
        yield run


def run_identity(run, mode):
    return {"id": run.id, "project": run.project, "entity": run.entity, "mode": mode} if run else None


def log_best(run, model, directory, epoch, nll):
    """Log inference weights and their selection result, excluding optimizer/data state."""
    import wandb
    path = directory/"best-decision.pt"
    save_checkpoint(path, model)
    artifact = wandb.Artifact(f"paya-{VERSION}-{run.id}", type="model",
                              metadata={"epoch": epoch, "selection_nll": nll, "version": VERSION})
    artifact.add_file(str(path), name="decision.pt")
    run.log_artifact(artifact, aliases=["best"])
    run.summary.update({"best/epoch": epoch, "best/selection_nll": nll})