from __future__ import annotations

import atexit
import threading
import time

import httpx


class Run:
    """Represents an active experiment run. Buffers metrics and flushes them periodically."""

    def __init__(
        self,
        run_id: int,
        project_id: int,
        api_url: str,
        api_token: str | None = None,
        flush_interval: float = 5.0,
    ):
        self._run_id = run_id
        self._project_id = project_id
        self._api_url = api_url.rstrip("/")
        self._api_token = api_token
        self._flush_interval = flush_interval

        self._metric_buffer: list[dict] = []
        self._param_buffer: dict[str, str | int | float | bool] = {}
        self._lock = threading.Lock()
        self._stopped = False

        # Background flush and heartbeat thread
        self._thread = threading.Thread(target=self._background_loop, daemon=True)
        self._thread.start()
        atexit.register(self.stop)

    @property
    def id(self) -> int:
        return self._run_id

    @classmethod
    def _create(
        cls,
        project: str,
        api_url: str,
        name: str | None = None,
        api_token: str | None = None,
    ) -> Run:
        """Create a run by calling the backend API."""
        # Resolve project name to project_id
        client = httpx.Client(base_url=api_url)
        headers = {"Authorization": f"Bearer {api_token}"} if api_token else {}

        # Try to find existing project or create one
        resp = client.get("/projects", headers=headers)
        resp.raise_for_status()
        projects = resp.json()["items"]
        project_name = project.split("/")[-1]  # Handle "workspace/project" format
        project_id = None
        for p in projects:
            if p["name"] == project_name:
                project_id = p["id"]
                break

        if project_id is None:
            resp = client.post("/projects", json={"name": project_name}, headers=headers)
            resp.raise_for_status()
            project_id = resp.json()["id"]

        # Init run
        resp = client.post(
            "/runs/init",
            json={"project_id": project_id, "name": name},
            headers=headers,
        )
        resp.raise_for_status()
        run_data = resp.json()
        client.close()

        return cls(
            run_id=run_data["id"],
            project_id=project_id,
            api_url=api_url,
            api_token=api_token,
        )

    def __setitem__(self, key: str, value: dict) -> None:
        """Set run parameters via dict syntax: run["parameters"] = {...}"""
        if isinstance(value, dict):
            with self._lock:
                self._param_buffer.update(value)
        else:
            raise TypeError("Value must be a dict of parameters")

    def log(self, name: str, value: float, step: int | None = None) -> None:
        """Log a metric value."""
        with self._lock:
            self._metric_buffer.append({"name": name, "value": value, "step": step})

    def _flush(self) -> None:
        """Flush buffered metrics and params to the backend."""
        with self._lock:
            metrics = self._metric_buffer[:]
            self._metric_buffer.clear()
            params = self._param_buffer.copy()
            self._param_buffer.clear()

        if not metrics and not params:
            return

        payload = {}
        if metrics:
            payload["metrics"] = metrics
        if params:
            payload["params"] = params

        try:
            with httpx.Client(base_url=self._api_url) as client:
                headers = (
                    {"Authorization": f"Bearer {self._api_token}"} if self._api_token else {}
                )
                client.post(f"/runs/{self._run_id}/log", json=payload, headers=headers)
        except Exception:
            # Re-add failed items to buffer
            with self._lock:
                self._metric_buffer = metrics + self._metric_buffer
                self._param_buffer = {**params, **self._param_buffer}

    def _heartbeat(self) -> None:
        """Send heartbeat to keep run alive."""
        try:
            with httpx.Client(base_url=self._api_url) as client:
                headers = (
                    {"Authorization": f"Bearer {self._api_token}"} if self._api_token else {}
                )
                client.post(f"/runs/{self._run_id}/heartbeat", headers=headers)
        except Exception:
            pass

    def _background_loop(self) -> None:
        """Periodically flush metrics and send heartbeats."""
        while not self._stopped:
            time.sleep(self._flush_interval)
            if self._stopped:
                break
            self._flush()
            self._heartbeat()

    def stop(self) -> None:
        """Finalize the run: flush remaining data and set status to completed."""
        if self._stopped:
            return
        self._stopped = True
        self._flush()

        try:
            with httpx.Client(base_url=self._api_url) as client:
                headers = (
                    {"Authorization": f"Bearer {self._api_token}"} if self._api_token else {}
                )
                client.patch(
                    f"/runs/{self._run_id}",
                    json={"status": "completed"},
                    headers=headers,
                )
        except Exception:
            pass
