from __future__ import annotations

import os
from typing import Any

from zvondb.client import RegistryClient
from zvondb.models import ArtifactInfo, JobDetailInfo, ProjectInfo


class Job:
    """Handle to a running job. Primary interface for resolve/create operations.

    Usage:
        job = exp.start_job(name="training", git_sha="abc123")
        emb = job.resolve("EmbeddingSet")
        path = job.create("SemanticIDSet", storage_name="my_s3")
        job.complete()

    Context manager usage:
        with exp.start_job(...) as job:
            emb = job.resolve("EmbeddingSet")
            path = job.create("Output", storage_name="my_s3")
            job.complete()
        # If an exception occurs before complete(), fail() is called automatically.

    W&B integration:
        with exp.start_job(name="train", git_sha="abc") as job:
            job.inject_wandb_env()
            # Sets WANDB_API_KEY, WANDB_ENTITY, WANDB_PROJECT, WANDB_RUN_ID
            # and auto-sets the tracking URL on the job.
    """

    def __init__(
        self,
        client: RegistryClient,
        project: str,
        experiment: str,
        job_id: str,
        project_info: ProjectInfo | None = None,
    ):
        self._client = client
        self._project = project
        self._experiment = experiment
        self._job_id = job_id
        self._project_info = project_info
        self._completed = False

    @property
    def id(self) -> str:
        return self._job_id

    def wandb_env(self) -> dict[str, str]:
        """Return W&B env vars for this job and auto-set the tracking URL.

        Returns WANDB_API_KEY, WANDB_ENTITY (from project config),
        WANDB_PROJECT (zvondb project name), and WANDB_RUN_ID (job ID).
        Also sets the tracking URL on the job so it's clickable in the UI.
        """
        env: dict[str, str] = {}
        if self._project_info:
            if self._project_info.wandb_api_key:
                env["WANDB_API_KEY"] = self._project_info.wandb_api_key
            if self._project_info.wandb_entity:
                env["WANDB_ENTITY"] = self._project_info.wandb_entity
        env["WANDB_PROJECT"] = self._project
        env["WANDB_RUN_ID"] = self._job_id

        entity = env.get("WANDB_ENTITY")
        if entity:
            url = f"https://wandb.ai/{entity}/{self._project}/runs/{self._job_id}"
            self.set_tracking_url(url)

        return env

    def inject_wandb_env(self) -> None:
        """Set W&B env vars in os.environ and auto-set the tracking URL.

        Calls wandb_env() and writes the result into os.environ.

        Frameworks like Axolotl that respect W&B env vars will automatically
        log to the correct project/run with no additional code.
        """
        for key, value in self.wandb_env().items():
            os.environ[key] = value

    def resolve(
        self, artifact_name: str, version: int | None = None
    ) -> ArtifactInfo:
        """Resolve an artifact via parent chain walk. Records input edge on this job."""
        data = self._client.get(
            f"/projects/{self._project}/experiments/{self._experiment}/resolve",
            params={
                "name": artifact_name,
                "job_id": self._job_id,
                "version": version,
            },
        )
        return ArtifactInfo(**data)

    def create(
        self,
        name: str,
        storage_name: str | None = None,
        description: str = "",
        metadata: dict[str, Any] | None = None,
        tags: list[str] | None = None,
        path: str | None = None,
    ) -> str:
        """Create an artifact. Returns the path.

        Either ``storage_name`` or ``path`` must be provided.

        - ``storage_name``: auto-generates the path from the project's storage backend.
        - ``path``: registers an existing path as-is (no storage backend needed).

        For v0 artifacts, use description to describe what the artifact is.
        For subsequent versions, use it as a changelog entry.
        """
        body: dict[str, Any] = {
            "name": name,
            "description": description,
            "metadata": metadata or {},
            "tags": tags or [],
            "job_id": self._job_id,
        }
        if storage_name is not None:
            body["storage_name"] = storage_name
        if path is not None:
            body["path"] = path
        data = self._client.post(
            f"/projects/{self._project}/experiments/{self._experiment}/artifacts",
            json=body,
        )
        return data["path"]

    def complete(self) -> None:
        """Mark this job as completed."""
        self._client.patch(
            f"/projects/{self._project}/jobs/{self._job_id}",
            json={"status": "completed"},
        )
        self._completed = True

    def fail(self, reason: str) -> None:
        """Mark this job as failed."""
        self._client.patch(
            f"/projects/{self._project}/jobs/{self._job_id}",
            json={"status": "failed", "failure_reason": reason},
        )
        self._completed = True

    def set_tracking_url(self, url: str) -> None:
        """Set a tracking URL (e.g. W&B dashboard) for this job."""
        self._client.patch(
            f"/projects/{self._project}/jobs/{self._job_id}",
            json={"tracking_url": url},
        )

    def detail(self) -> JobDetailInfo:
        """Get full job detail including inputs and outputs."""
        data = self._client.get(
            f"/projects/{self._project}/jobs/{self._job_id}"
        )
        return JobDetailInfo(**data)

    def __enter__(self) -> Job:
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        if exc_type is not None and not self._completed:
            self.fail(reason=f"{exc_type.__name__}: {exc_val}")
