from __future__ import annotations

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


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

    def create_project(
        self,
        storage_backends: dict[str, str] | None = None,
        wandb_api_key: str | None = None,
        wandb_entity: str | None = None,
    ) -> ProjectInfo:
        payload: dict = {
            "name": self._project,
            "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 = self._client.post("/projects", json=payload)
        info = ProjectInfo(**data)
        self._project_info = info
        return info

    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, 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 = "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 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()
