from __future__ import annotations

from zvondb.client import RegistryClient
from zvondb.exceptions import NotFoundError
from zvondb.experiment import Experiment
from zvondb.models import ExperimentInfo, MetricsQueryRow, ProjectInfo, UniversalArtifactInfo


class Registry:
    """Entry point for the ML Artifact Registry SDK.

    Usage:
        registry = Registry(project="semantic-id")
        exp = registry.get_experiment("rqvae-exp-1")
        job = exp.start_job(name="training", git_sha="abc123")
    """

    def __init__(
        self,
        project: str,
        base_url: str | None = None,
        api_key: str | None = None,
    ):
        self._client = RegistryClient(base_url=base_url, api_key=api_key)
        self._project = project
        self._project_info: ProjectInfo | None = None

    @property
    def project(self) -> str:
        return self._project

    def _get_project_info(self) -> ProjectInfo:
        if self._project_info is None:
            self._project_info = self.get_project()
        return self._project_info

    @staticmethod
    def create_project(
        name: str,
        storage_backends: dict[str, str] | None = None,
        wandb_api_key: str | None = None,
        wandb_entity: str | None = None,
        base_url: str | None = None,
        api_key: str | None = None,
    ) -> ProjectInfo:
        """Create a new project. This is a one-time setup operation.

        Usage:
            Registry.create_project(
                "my-project",
                storage_backends={"s3": "s3://bucket/prefix/"},
            )
        """
        client = RegistryClient(base_url=base_url, api_key=api_key)
        payload: dict = {
            "name": name,
            "storage_backends": storage_backends or {},
        }
        if wandb_api_key:
            payload["wandb_api_key"] = wandb_api_key
        if wandb_entity:
            payload["wandb_entity"] = wandb_entity
        data = client.post("/projects", json=payload)
        client.close()
        return ProjectInfo(**data)

    def get_project(self) -> ProjectInfo:
        data = self._client.get(f"/projects/{self._project}")
        info = ProjectInfo(**data)
        self._project_info = info
        return info

    def get_experiment(self, name: str) -> Experiment:
        project_info = self._get_project_info()
        # Validate experiment exists
        self._client.get(f"/projects/{self._project}/experiments/{name}")
        return Experiment(self._client, self._project, name, project_info=project_info)

    def create_experiment(
        self, name: str, parent: str | None = "baseline", description: str = ""
    ) -> Experiment:
        project_info = self._get_project_info()
        self._client.post(
            f"/projects/{self._project}/experiments",
            json={"name": name, "parent": parent, "description": description},
        )
        return Experiment(self._client, self._project, name, project_info=project_info)

    def ensure_experiment(
        self, name: str, parent: str | None = "baseline", description: str = ""
    ) -> Experiment:
        """Get experiment if exists, create if missing."""
        project_info = self._get_project_info()
        try:
            self._client.get(f"/projects/{self._project}/experiments/{name}")
        except NotFoundError:
            self._client.post(
                f"/projects/{self._project}/experiments",
                json={"name": name, "parent": parent, "description": description},
            )
        return Experiment(self._client, self._project, name, project_info=project_info)

    def query_metrics(
        self,
        job_name: str,
        experiments: list[str] | None = None,
        keys: list[str] | None = None,
    ) -> list[MetricsQueryRow]:
        """Query metrics across experiments for completed jobs.

        Args:
            job_name: Filter by job name (required).
            experiments: Optional list of experiment names to filter by.
            keys: Optional list of metric keys to include.

        Returns:
            List of MetricsQueryRow with experiment, job info, and metrics.
        """
        params: dict[str, str] = {"job_name": job_name}
        if experiments:
            params["experiments"] = ",".join(experiments)
        if keys:
            params["keys"] = ",".join(keys)
        data = self._client.get(
            f"/projects/{self._project}/metrics", params=params
        )
        return [MetricsQueryRow(**row) for row in data["rows"]]

    def list_universal(
        self, type: str | None = None
    ) -> list[UniversalArtifactInfo]:
        """List universal artifacts, optionally filtered by type."""
        params: dict[str, str] = {}
        if type:
            params["type"] = type
        data = self._client.get("/universal/artifacts", params=params)
        return [UniversalArtifactInfo(**item) for item in data]

    def list_universal_types(self) -> list[str]:
        """List distinct artifact types in the universal registry."""
        return self._client.get("/universal/artifacts/types")

    def promote_universal(
        self,
        name: str,
        type: str,
        artifact_id: str,
        description: str = "",
        metadata: dict | None = None,
        tags: list[str] | None = None,
    ) -> UniversalArtifactInfo:
        """Promote an artifact to the universal registry.

        If a universal artifact with the same name+type already exists,
        it is updated to point to the new artifact_id.
        """
        data = self._client.post(
            "/universal/artifacts",
            json={
                "name": name,
                "type": type,
                "artifact_id": artifact_id,
                "description": description,
                "metadata": metadata or {},
                "tags": tags or [],
            },
        )
        return UniversalArtifactInfo(**data)

    def list_experiments(self) -> list[ExperimentInfo]:
        data = self._client.get(f"/projects/{self._project}/experiments")
        return [ExperimentInfo(**e) for e in data]

    def close(self) -> None:
        self._client.close()

    def __enter__(self) -> Registry:
        return self

    def __exit__(self, *args) -> None:
        self.close()
