from __future__ import annotations

import os
from typing import Any

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

_FIELD_LIMITS = {
    "name": 255,
    "type": 255,
    "storage_name": 255,
    "path": 2048,
}


def _validate_lengths(**kwargs: str | None) -> None:
    for field, value in kwargs.items():
        if value is None:
            continue
        limit = _FIELD_LIMITS.get(field)
        if limit and len(value) > limit:
            raise ValueError(
                f"'{field}' is {len(value)} chars, max is {limit}: "
                f"{value[:80]}{'...' if len(value) > 80 else ''}"
            )


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.

    Dry-run mode (report=False):
        with exp.start_job(name="debug", git_sha="abc", report=False) as job:
            data = job.resolve("TrainingData")  # resolves but records nothing
            # job.create(...) would raise RuntimeError
            job.complete()  # no-op

    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 | None = None,
        project_info: ProjectInfo | None = None,
        report: bool = True,
    ):
        self._client = client
        self._project = project
        self._experiment = experiment
        self._job_id = job_id
        self._project_info = project_info
        self._report = report
        self._completed = False

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

    @property
    def report(self) -> bool:
        return self._report

    def _require_report(self, method: str) -> None:
        if not self._report:
            raise RuntimeError(
                f"{method}() is not available in dry-run mode (report=False)"
            )

    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.
        """
        self._require_report("wandb_env")
        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.

        In normal mode, records an input edge on this job.
        In dry-run mode (report=False), resolves without recording anything.
        """
        params: dict[str, Any] = {
            "name": artifact_name,
            "version": version,
        }
        if self._report:
            params["job_id"] = self._job_id
        data = self._client.get(
            f"/projects/{self._project}/experiments/{self._experiment}/resolve",
            params=params,
        )
        return ArtifactInfo(**data)

    def create(
        self,
        name: str,
        type: 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.

        Not available in dry-run mode (report=False).
        """
        self._require_report("create")
        _validate_lengths(name=name, type=type, storage_name=storage_name, path=path)
        if storage_name is None and path is None:
            raise ValueError("Either 'storage_name' or 'path' must be provided")
        body: dict[str, Any] = {
            "name": name,
            "type": type,
            "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 create_with_info(
        self,
        name: str,
        type: str,
        storage_name: str | None = None,
        description: str = "",
        metadata: dict[str, Any] | None = None,
        tags: list[str] | None = None,
        path: str | None = None,
    ) -> ArtifactInfo:
        """Like create(), but returns the full ArtifactInfo instead of just the path.

        Not available in dry-run mode (report=False).
        """
        self._require_report("create_with_info")
        _validate_lengths(name=name, type=type, storage_name=storage_name, path=path)
        if storage_name is None and path is None:
            raise ValueError("Either 'storage_name' or 'path' must be provided")
        body: dict[str, Any] = {
            "name": name,
            "type": type,
            "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 ArtifactInfo(**data)

    def complete(self) -> None:
        """Mark this job as completed. No-op in dry-run mode."""
        if not self._report:
            self._completed = True
            return
        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. No-op in dry-run mode."""
        if not self._report:
            self._completed = True
            return
        self._client.patch(
            f"/projects/{self._project}/jobs/{self._job_id}",
            json={"status": "failed", "failure_reason": reason},
        )
        self._completed = True

    def resolve_universal(self, name: str, type: str) -> ArtifactInfo:
        """Resolve an artifact from the universal registry.

        Records an input edge on this job (with source='universal'),
        so the lineage graph shows this dependency came from the universal registry.

        Not available in dry-run mode (report=False).
        """
        self._require_report("resolve_universal")
        params = {"name": name, "type": type, "job_id": self._job_id}
        data = self._client.get("/universal/artifacts/resolve", params=params)
        return ArtifactInfo(**data)

    def log_metrics(self, metrics: dict[str, float]) -> None:
        """Log key-value metrics for this job. Can be called multiple times — keys are upserted.

        Not available in dry-run mode (report=False).
        """
        self._require_report("log_metrics")
        self._client.post(
            f"/projects/{self._project}/jobs/{self._job_id}/metrics",
            json={"metrics": metrics},
        )

    def set_tracking_url(self, url: str) -> None:
        """Set a tracking URL (e.g. W&B dashboard) for this job."""
        self._require_report("set_tracking_url")
        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."""
        self._require_report("detail")
        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}")
