"""Tyko client for API communication."""

import atexit
import hashlib
import multiprocessing
import os
import time
from datetime import datetime
from typing import Any, Mapping

import httpx

from ._environment import capture_environment
from ._ipc import IPCConnection, _socket_path
from ._sidecar import SidecarConfig, run_sidecar
from ._types import Experiment, Project, Run, RunStatus


def _make_client_id(*, api_url: str, api_key: str | None) -> str:
    """Derive a deterministic client ID from the process PID and API config.

    Same process + same config always produces the same ID, so re-creating
    TykoClient() in a notebook re-run finds and reuses the existing sidecar
    rather than spawning a new orphaned process.
    """
    raw = f"{os.getpid()}-{api_url}-{api_key or ''}"
    return hashlib.sha1(raw.encode()).hexdigest()[:16]


class TykoError(Exception):
    """Base exception for all Tyko SDK errors."""

    def __init__(self, message: str, *, status_code: int | None = None) -> None:
        self.status_code = status_code
        super().__init__(message)


class AuthenticationError(TykoError):
    """Raised when authentication fails (invalid or missing API key)."""


def _raise_for_status(response: httpx.Response) -> None:
    """Check response status and raise a clear TykoError on failure."""
    if response.is_success:
        return

    # Try to extract detail from Litestar's JSON error response
    detail = None
    try:
        body = response.json()
        detail = body.get("detail") or body.get("message")
    except Exception:
        pass

    status = response.status_code

    if status == 401:
        msg = detail or "Authentication failed"
        raise AuthenticationError(
            f"{msg}. Check your API key (TYKO_API_KEY env var or api_key= parameter).",
            status_code=status,
        )

    if status == 403:
        raise TykoError(
            detail or "Permission denied.",
            status_code=status,
        )

    if status == 404:
        raise TykoError(
            detail or "Resource not found.",
            status_code=status,
        )

    if status >= 500:
        raise TykoError(
            f"Server error ({status}): {detail or 'An internal error occurred. Please try again.'}",
            status_code=status,
        )

    raise TykoError(
        f"Request failed ({status}): {detail or response.reason_phrase}",
        status_code=status,
    )


class TykoClient:
    """Client for interacting with the Tyko Labs API."""

    _ipc: IPCConnection | None

    def __init__(self, api_url: str | None = None, api_key: str | None = None):
        """Initialize the Tyko client.

        Spawns a sidecar process that handles all HTTP communication and
        system monitoring. The sidecar is started in the background — if it
        fails to start, all operations fall back to synchronous HTTP.

        Args:
            api_url: The base URL for the Tyko API. Defaults to the production
                server at https://api.tyko-labs.com.
            api_key: API key for authentication. If not provided, will check
                the TYKO_API_KEY environment variable. If still not found,
                an anonymous session will be created automatically.
        """

        if api_url is None:
            api_url = os.environ.get("TYKO_API_URL", "https://api.tyko-labs.com")

        self.api_url = api_url.rstrip("/")

        # Get API key from argument or environment variable.
        # When None, requests are sent without auth (works with ALLOW_ANONYMOUS servers).
        self.api_key = api_key or os.environ.get("TYKO_API_KEY")

        self._client_uuid = _make_client_id(api_url=self.api_url, api_key=self.api_key)
        self._ipc = None
        # TYKO_DISABLE_SIDECAR disables the sidecar (used in tests)
        if not os.environ.get("TYKO_DISABLE_SIDECAR"):
            self._start_sidecar()

    def _start_sidecar(self) -> None:
        """Spawn the sidecar process and connect to it.

        On any failure, self._ipc stays None and all calls fall back to
        synchronous HTTP (identical to the pre-sidecar behavior).
        """
        sock_path = _socket_path(client_uuid=self._client_uuid)
        config = SidecarConfig(
            client_uuid=self._client_uuid,
            api_url=self.api_url,
            api_key=self.api_key,
            main_pid=os.getpid(),
            socket_path=sock_path,
        )

        # Try to reconnect to an existing sidecar (notebook restart safety)
        try:
            conn = IPCConnection.connect(client_uuid=self._client_uuid, timeout=0.3)
            self._ipc = conn
            return
        except OSError:
            pass

        # Spawn a fresh sidecar process
        try:
            proc = multiprocessing.Process(
                target=run_sidecar,
                kwargs={"config": config},
                daemon=True,
            )
            proc.start()
        except Exception:
            return

        # Poll for the socket to appear (up to 5 seconds)
        deadline = time.monotonic() + 5.0
        while time.monotonic() < deadline:
            try:
                conn = IPCConnection.connect(client_uuid=self._client_uuid, timeout=0.3)
                self._ipc = conn
                return
            except OSError:
                time.sleep(0.05)

        # Sidecar failed to start in time — fall back to sync HTTP

    def _headers(self) -> dict[str, str]:
        headers: dict[str, str] = {}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"
        return headers

    def _post(self, path: str, data: Mapping[str, Any]) -> httpx.Response:
        url = f"{self.api_url}{path}"
        response = httpx.post(url, json=data, headers=self._headers())
        _raise_for_status(response)
        return response

    def _patch(self, path: str, data: Mapping[str, Any]) -> httpx.Response:
        url = f"{self.api_url}{path}"
        response = httpx.patch(url, json=data, headers=self._headers())
        _raise_for_status(response)
        return response

    def update_status(self, run_id: str, status: RunStatus) -> None:
        """Update the status of a run.

        Args:
            run_id: The ID of the run to update.
            status: The new status for the run.
        """
        self._patch(f"/runs/{run_id}", {"status": status})

    def close_run(self, run_id: str, status: RunStatus, at: datetime) -> None:
        """Close a run with the given status and timestamp.

        Args:
            run_id: The ID of the run to close.
            status: The final status of the run ('completed' or 'failed').
            at: Timestamp of when the run ended.
        """
        self._patch(f"/runs/{run_id}", {"status": status, "at": at.isoformat()})

    def update_params(self, run_id: str, params: dict[str, object]) -> None:
        """Update the parameters for a run.

        Params values are merged with existing params (new keys are added,
        existing keys are overwritten).

        Args:
            run_id: The ID of the run to update.
            params: Dictionary of parameter values to set/update.
        """
        self._patch(f"/runs/{run_id}", {"params": params})

    def update_environment(self, run_id: str, environment: dict[str, object]) -> None:
        """Update the system information for a run.

        System info values are merged with existing system info (new keys are added,
        existing keys are overwritten).

        Args:
            run_id: The ID of the run to update.
            environment: Dictionary of system information to set/update.
        """
        self._patch(f"/runs/{run_id}", {"environment": environment})

    def log_metrics(
        self,
        run_id: str,
        values: dict[str, float | int],
        *,
        _at: int,
        _relative_time_ns: int,
    ) -> None:
        """Log metric values for a run.

        Each call creates a new entry with a sequential number auto-assigned
        by the server for ordering.

        Metric names can use prefixes for grouping in the UI:
        - "train/loss", "train/accuracy" → grouped under "train"
        - "eval/loss", "eval/accuracy" → grouped under "eval"
        - Underscore also works: "train_loss" → grouped under "train"

        Args:
            run_id: The ID of the run to log metrics for.
            values: Dictionary of metric names to values.
            _at: Monotonic nanosecond timestamp from the client for ordering.
            _relative_time_ns: Nanoseconds since run started (for aligning metrics).
        """
        entry = {
            "timestamp": _at,
            "relative_time_ns": _relative_time_ns,
            "values": values,
        }
        data = {"entries": [entry]}
        self._post(f"/runs/{run_id}/entries", data)

    def start_run(
        self,
        project: str,
        experiment: str | None = None,
        params: dict[str, object] | None = None,
        *,
        system_monitor_interval: float | None = 15.0,
        log_system: bool = True,
        capture_output: bool = True,
    ) -> Run:
        """Start a new run.

        This is the main method for tracking experiments. It creates a run
        within the specified project and optional experiment.

        Environment information (Python version, platform, CPU, RAM, GPU) is
        automatically captured and sent to the server. System metrics (CPU %,
        RAM, GPU utilization) are sampled periodically by the sidecar process.

        Args:
            project: The project name. Required.
            experiment: Optional experiment name. If not provided, the server
                will use the default experiment.
            params: Optional dictionary of parameters to log with the run.
                These are static values like hyperparameters that don't change
                during the run.
            system_monitor_interval: Seconds between system metric samples.
                Defaults to 15 seconds. Set to None to disable.
            log_system: Whether to log system metrics. Set to False for
                non-zero DDP ranks to avoid duplicate system metrics.

        Returns:
            A Run instance that can be used as a context manager.
        """
        data: dict[str, Any] = {"project": project}
        if experiment is not None:
            data["experiment"] = experiment
        if params is not None:
            data["params"] = params

        # Auto-capture environment information
        data["environment"] = capture_environment()

        response = self._post("/runs", data)

        response_json = response.json()

        run_id = response_json["id"]
        run_name = response_json["name"]
        project_name = response_json["project"]["name"]
        experiment_name = response_json["experiment"]["name"]
        sequential = response_json["sequential"]

        run = Run(
            id=run_id,
            name=run_name,
            sequential=sequential,
            experiment=Experiment(name=experiment_name),
            project=Project(name=project_name),
            _client=self,
            _client_uuid=self._client_uuid,
            _capture_output=capture_output,
        )

        # Wire the IPC connection into the run (if sidecar is available)
        effective_interval = system_monitor_interval if log_system else None
        if self._ipc is not None:
            ok = self._ipc.send_and_wait(
                {
                    "type": "set_run",
                    "run_id": run_id,
                    "system_monitor_interval": effective_interval,
                }
            )
            if ok:
                run._ipc[0] = self._ipc
                run.params._ipc = self._ipc
                run.environment._ipc = self._ipc

        # Ensure the run is marked completed/failed even if __exit__ is never called
        # (e.g. training script exits without using the context manager)
        atexit.register(run.__exit__, None, None, None)

        return run

    def close(self) -> None:
        """Shut down the sidecar process and release resources.

        Call this when you are done using the client and no run is active.
        If using start_run() as a context manager, the sidecar is shut down
        automatically on __exit__.
        """
        if self._ipc is not None:
            self._ipc.send_and_wait({"type": "shutdown"})
            self._ipc.close()
            self._ipc = None
