"""Interactive GPU sessions with device-like interface."""
from __future__ import annotations

import time
from pathlib import Path
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from .client import Client


class Session:
    """GPU session with a device-like interface.

    Use as a context manager::

        with yarn.session(gpu="rtx4090") as device:
            model = MyModel().to(device)
            device.log({"epoch": 1, "loss": 0.42})
            device.save("model.pt", model.state_dict())
    """

    def __init__(self, data: dict, client: Client):
        self._data = data
        self._client = client
        self._device_str: str | None = None

    # -- Properties ------------------------------------------------------------

    @property
    def id(self) -> str:
        return self._data.get("id", "")

    @property
    def name(self) -> str:
        return self._data.get("session_name", self._data.get("name", ""))

    @property
    def status(self) -> str:
        return self._data.get("status", "")

    @property
    def ray_address(self) -> str:
        return self._data.get("ray_endpoint", "")

    @property
    def dashboard_url(self) -> str:
        return self._data.get("dashboard_url", "")

    @property
    def gpu_type(self) -> str:
        return self._data.get("gpu_type", "")

    @property
    def created_at(self) -> str:
        return self._data.get("created_at", "")

    # -- Device interface ------------------------------------------------------

    @property
    def type(self) -> str:
        """PyTorch-compatible device type string (e.g. 'cuda')."""
        return "cuda"

    @property
    def index(self) -> int:
        """Device index (always 0 for single-GPU sessions)."""
        return 0

    def __str__(self) -> str:
        """Return 'cuda:0' — passable to tensor.to() when running inside the session."""
        return "cuda:0"

    def __repr__(self) -> str:
        return f"Session({self.name!r}, gpu={self.gpu_type!r}, status={self.status!r})"

    # -- Lifecycle -------------------------------------------------------------

    def refresh(self) -> None:
        """Re-fetch session state from the API."""
        self._data = self._client.get(f"/v1/research/sessions/{self.id}")

    def wait_until_ready(
        self, timeout: float = 300, poll_interval: float = 5.0
    ) -> None:
        """Block until the session is running or has failed."""
        deadline = time.time() + timeout
        while time.time() < deadline:
            self.refresh()
            if self.status.lower() == "running":
                return
            if self.status.lower() == "failed":
                raise RuntimeError(f"Session {self.id} failed")
            time.sleep(min(poll_interval, max(0, deadline - time.time())))
        raise TimeoutError(
            f"Session {self.id} not ready within {timeout}s (status: {self.status})"
        )

    def stop(self) -> None:
        """Delete the session."""
        self._client.delete(f"/v1/research/sessions/{self.id}")

    def __enter__(self) -> Session:
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        self.stop()

    # -- Logging + artifacts ---------------------------------------------------

    def log(self, metrics: dict[str, Any], step: int | None = None) -> None:
        """Log metrics for this session.

        Args:
            metrics: Key-value pairs (e.g. {"loss": 0.42, "epoch": 1}).
            step: Optional training step number.
        """
        body: dict[str, Any] = {"metrics": metrics}
        if step is not None:
            body["step"] = step
        try:
            self._client.post(f"/v1/research/sessions/{self.id}/metrics", json=body)
        except Exception:
            pass  # Best-effort — don't break training on metrics failure

    def save(self, key: str, data: Any = None, *, local_path: str | None = None) -> dict:
        """Save an artifact (checkpoint, model weights, etc.).

        Either provide ``data`` (bytes or dict that will be serialized) or
        ``local_path`` (path to a file to upload).

        Args:
            key: Storage key / filename (e.g. "model.pt", "results.json").
            data: Raw bytes or a dict (serialized as JSON).
            local_path: Path to a local file to upload instead of data.

        Returns:
            Upload response dict.
        """
        storage_key = f"sessions/{self.id}/{key}"
        if local_path:
            return self._client.storage.upload(storage_key, local_path)
        if isinstance(data, (bytes, bytearray)):
            import tempfile, os
            tmp = tempfile.NamedTemporaryFile(delete=False, suffix=Path(key).suffix)
            try:
                tmp.write(data)
                tmp.close()
                return self._client.storage.upload(storage_key, tmp.name)
            finally:
                os.unlink(tmp.name)
        if isinstance(data, dict):
            import json, tempfile, os
            tmp = tempfile.NamedTemporaryFile(
                delete=False, suffix=".json", mode="w"
            )
            try:
                json.dump(data, tmp)
                tmp.close()
                return self._client.storage.upload(storage_key, tmp.name)
            finally:
                os.unlink(tmp.name)
        raise TypeError(f"data must be bytes or dict, got {type(data).__name__}")


class Sessions:
    def __init__(self, client: Client):
        self._client = client

    def create(
        self,
        name: str,
        *,
        gpu: str = "rtx-4090",
        gpu_count: int = 1,
        cpu: str = "4",
        memory: str = "8Gi",
        image: str = "",
        idle_timeout_minutes: int = 60,
        env_vars: dict[str, str] | None = None,
        compute_target: str | None = None,
        provider: str | None = None,
        region: str | None = None,
        sovereignty: bool = False,
        min_vram: int | None = None,
    ) -> Session:
        """Create an interactive GPU session."""
        body: dict = {
            "name": name,
            "gpu_type": gpu,
            "gpu_count": gpu_count,
            "cpu": cpu,
            "memory": memory,
            "idle_timeout_minutes": idle_timeout_minutes,
        }
        if image:
            body["image"] = image
        if env_vars:
            body["env_vars"] = env_vars
        if compute_target:
            body["compute_target"] = compute_target
        if provider:
            body["provider"] = provider
        if region:
            body["region"] = region
        if sovereignty:
            body["sovereignty"] = True
        if min_vram:
            body["min_vram"] = min_vram
        data = self._client.post("/v1/research/sessions", json=body)
        return Session(data, self._client)

    def get(self, session_id: str) -> Session:
        """Get session details."""
        data = self._client.get(f"/v1/research/sessions/{session_id}")
        return Session(data, self._client)

    def delete(self, session_id: str) -> None:
        """Delete a session."""
        self._client.delete(f"/v1/research/sessions/{session_id}")

    def list(self) -> list[dict]:
        """List all sessions."""
        return self._client.get("/v1/research/sessions")

    def session(
        self,
        name: str,
        *,
        gpu: str = "rtx-4090",
        gpu_count: int = 1,
        cpu: str = "4",
        memory: str = "8Gi",
        image: str = "",
        idle_timeout_minutes: int = 60,
        env_vars: dict[str, str] | None = None,
        compute_target: str | None = None,
        provider: str | None = None,
        region: str | None = None,
        sovereignty: bool = False,
        min_vram: int | None = None,
        wait_timeout: float = 300,
    ) -> Session:
        """Create a session and wait until it is ready.

        Returns a ``Session`` that can be used as a context manager::

            with client.sessions.session("exp-1") as device:
                model = MyModel().to(device)
                device.log({"loss": 0.42})
        """
        s = self.create(
            name,
            gpu=gpu,
            gpu_count=gpu_count,
            cpu=cpu,
            memory=memory,
            image=image,
            idle_timeout_minutes=idle_timeout_minutes,
            env_vars=env_vars,
            compute_target=compute_target,
            provider=provider,
            region=region,
            sovereignty=sovereignty,
            min_vram=min_vram,
        )
        s.wait_until_ready(timeout=wait_timeout)
        return s
