from collections import UserDict
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import StrEnum
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from .client import TykoClient


class RunStatus(StrEnum):
    COMPLETED = "completed"
    FAILED = "failed"


@dataclass(frozen=True, slots=True)
class Experiment:
    name: str


@dataclass(frozen=True, slots=True)
class Project:
    name: str


class RunParams(UserDict[str, object]):
    """A dict-like object for run parameters that syncs to the server.

    Supports dict-style access for setting values. Each modification
    is immediately synced to the server.

    Example:
        >>> with client.start_run(project="my-project") as run:
        ...     run.params["learning_rate"] = 0.001
        ...     run.params["batch_size"] = 32
        ...     run.params.update({"model": "resnet50", "epochs": 100})
    """

    def __init__(self, run_id: str, client: "TykoClient") -> None:
        super().__init__()
        self._run_id = run_id
        self._client = client

    def __setitem__(self, key: str, value: object) -> None:
        self.data[key] = value
        self._client.update_params(self._run_id, {key: value})

    def __repr__(self) -> str:
        return f"RunParams({self.data!r})"

    def update(self, values: dict[str, object], **kwargs: object) -> None:  # type: ignore[override]
        self.data.update(values, **kwargs)
        self._client.update_params(self._run_id, {**values, **kwargs})


class Environment(UserDict[str, object]):
    """A dict-like object for system information that syncs to the server.

    Supports dict-style access for setting values. Each modification
    is immediately synced to the server.

    Example:
        >>> with client.start_run(project="my-project") as run:
        ...     run.environment["python_version"] = "3.12"
        ...     run.environment["gpu_model"] = "A100"
        ...     run.environment.update({"git_commit": "abc123", "cuda_version": "12.1"})
    """

    def __init__(self, run_id: str, client: "TykoClient") -> None:
        super().__init__()
        self._run_id = run_id
        self._client = client

    def __setitem__(self, key: str, value: object) -> None:
        self.data[key] = value
        self._client.update_environment(self._run_id, {key: value})

    def __repr__(self) -> str:
        return f"Environment({self.data!r})"

    def update(self, values: dict[str, object], **kwargs: object) -> None:  # type: ignore[override]
        self.data.update(values, **kwargs)
        self._client.update_environment(self._run_id, {**values, **kwargs})


@dataclass(frozen=True, slots=True)
class Run:
    id: str
    name: str
    sequential: int
    experiment: Experiment
    project: Project
    _client: "TykoClient" = field(repr=False)
    params: RunParams = field(init=False, repr=False)
    environment: Environment = field(init=False, repr=False)

    def __post_init__(self) -> None:
        object.__setattr__(self, "params", RunParams(self.id, self._client))
        object.__setattr__(self, "environment", Environment(self.id, self._client))

    @property
    def url(self) -> str:
        return f"https://tyko-labs.com/dashboard/projects/{self.project.name}/experiments/{self.experiment.name}/runs/{self.name}-{self.sequential}"

    def capture_environment(self) -> dict[str, object]:
        """Automatically capture and record environment information.

        Captures:
        - Python version
        - Platform/OS
        - CPU count
        - RAM size (if psutil available)
        - GPU count and names (if torch available)

        The captured information is automatically sent to the server.

        Returns:
            Dictionary of captured environment information
        """
        from .environment import capture_environment

        env_info = capture_environment()
        self.environment.update(env_info)
        return env_info

    def __enter__(self) -> "Run":
        return self

    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
        status = RunStatus.FAILED if exc_type is not None else RunStatus.COMPLETED
        at = datetime.now(timezone.utc)
        self._client.close_run(self.id, status, at)
