import logging
import os
import sys
import threading
import time
from collections import UserDict
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import StrEnum
from typing import TYPE_CHECKING, Any, TextIO

if TYPE_CHECKING:
    from ._client import TykoClient
    from ._ipc import IPCConnection


class RunStatus(StrEnum):
    CREATED = "created"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    STOPPED = "stopped"
    UNKNOWN = "unknown"


@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 (via IPC to the sidecar, or
    directly via HTTP if the sidecar is not available).

    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", ipc: "IPCConnection | None" = None) -> None:
        super().__init__()
        self._run_id = run_id
        self._client = client
        self._ipc: "IPCConnection | None" = ipc

    def __setitem__(self, key: str, value: object) -> None:
        self.data[key] = value
        if self._ipc is not None:
            self._ipc.send({"type": "params", "values": {key: value}})
        else:
            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)
        merged = {**values, **kwargs}
        if self._ipc is not None:
            self._ipc.send({"type": "params", "values": merged})
        else:
            self._client.update_params(self._run_id, merged)


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 (via IPC to the sidecar, or
    directly via HTTP if the sidecar is not available).

    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", ipc: "IPCConnection | None" = None) -> None:
        super().__init__()
        self._run_id = run_id
        self._client = client
        self._ipc: "IPCConnection | None" = ipc

    def __setitem__(self, key: str, value: object) -> None:
        self.data[key] = value
        if self._ipc is not None:
            self._ipc.send({"type": "environment", "values": {key: value}})
        else:
            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)
        merged = {**values, **kwargs}
        if self._ipc is not None:
            self._ipc.send({"type": "environment", "values": merged})
        else:
            self._client.update_environment(self._run_id, merged)


class _TeeWriter:
    """Wraps a text stream to write-through to the original AND send lines via IPC."""

    def __init__(
        self,
        original: TextIO,
        key: str,
        ipc: "IPCConnection",
        start_time_ns: int,
    ) -> None:
        self._original = original
        self._key = key
        self._ipc = ipc
        self._start_time_ns = start_time_ns
        self._buffer = ""
        self._lock = threading.Lock()

    def write(self, text: str) -> int:
        n = self._original.write(text)
        with self._lock:
            self._buffer += text
            while "\n" in self._buffer:
                line, self._buffer = self._buffer.split("\n", 1)
                if line:
                    self._send(line)
        return n

    def flush(self) -> None:
        self._original.flush()
        with self._lock:
            if self._buffer:
                self._send(self._buffer)
                self._buffer = ""

    def _send(self, line: str) -> None:
        now_ns = time.time_ns()
        self._ipc.send(
            {
                "type": "log",
                "values": {self._key: line},
                "timestamp": now_ns,
                "relative_time_ns": time.monotonic_ns() - self._start_time_ns,
            }
        )

    def fileno(self) -> int:
        return self._original.fileno()

    @property
    def buffer(self) -> Any:
        return getattr(self._original, "buffer", None)

    def __getattr__(self, name: str) -> Any:
        return getattr(self._original, name)


class TykoLogHandler(logging.Handler):
    """Logging handler that captures log records and sends them via IPC."""

    def __init__(self, *, ipc: "IPCConnection", start_time_ns: int) -> None:
        super().__init__()
        self._ipc = ipc
        self._start_time_ns = start_time_ns

    def emit(self, record: logging.LogRecord) -> None:
        try:
            line = self.format(record)
            now_ns = time.time_ns()
            self._ipc.send(
                {
                    "type": "log",
                    "values": {"console/stdout": line},
                    "timestamp": now_ns,
                    "relative_time_ns": time.monotonic_ns() - self._start_time_ns,
                }
            )
        except Exception:
            self.handleError(record)


@dataclass(frozen=True, slots=True)
class Run:
    id: str
    name: str
    sequential: int
    experiment: Experiment
    project: Project
    _client: "TykoClient" = field(repr=False)
    _client_uuid: str = field(repr=False)
    _start_time_ns: int = field(default_factory=time.monotonic_ns, repr=False)
    _capture_output: bool = field(default=True, repr=False)
    params: RunParams = field(init=False, repr=False)
    environment: Environment = field(init=False, repr=False)
    # Mutable container to track if run has started (for frozen dataclass)
    _is_running: list[bool] = field(init=False, repr=False)
    # Mutable container for IPC connection (set after sidecar connects)
    _ipc: "list[IPCConnection | None]" = field(init=False, repr=False)
    # Mutable containers for console capture (stdout tee, stderr tee, log handler)
    _tee_writers: "list[_TeeWriter | None]" = field(init=False, repr=False)
    _log_handler: "list[TykoLogHandler | None]" = field(init=False, repr=False)

    def __post_init__(self) -> None:
        object.__setattr__(self, "_is_running", [False])
        object.__setattr__(self, "_ipc", [None])
        object.__setattr__(self, "_tee_writers", [None, None])  # [stdout_tee, stderr_tee]
        object.__setattr__(self, "_log_handler", [None])
        object.__setattr__(self, "params", RunParams(run_id=self.id, client=self._client))
        object.__setattr__(self, "environment", Environment(run_id=self.id, client=self._client))

    def _get_ipc(self) -> "IPCConnection | None":
        return self._ipc[0]  # type: ignore[return-value]

    @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 should_stop(self) -> bool:
        """Check if a stop has been requested for this run.

        The sidecar polls GET /runs/{id}/commands every 5 seconds. When a stop
        command arrives, it writes a sentinel file that this method checks.

        Call this in your training loop to support graceful stop:

        Example:
            >>> with client.start_run(project="mnist") as run:
            ...     for epoch in range(1000):
            ...         if run.should_stop():
            ...             break
            ...         train_one_epoch()
            ...         run.log({"loss": loss})

        Returns:
            True if a stop has been requested, False otherwise.
        """
        from ._sidecar import stop_sentinel_path

        sentinel = stop_sentinel_path(client_uuid=self._client_uuid)
        return os.path.exists(sentinel)

    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 log(self, values: dict[str, float | int]) -> None:
        """Log metric values for this run.

        Each call creates a new entry. The server auto-assigns a sequential
        number for ordering, and the client provides both:
        - monotonic nanosecond timestamp for precise ordering
        - relative_time_ns: nanoseconds since run started (for scatter plots)

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

        This matches the convention used by HuggingFace Transformers and
        other ML frameworks.

        Note: If the run is not yet in "running" status, calling log() will
        automatically transition it to "running". This ensures the UI shows
        the correct status even if the context manager isn't used.

        Args:
            values: Dictionary of metric names to numeric values.

        Example:
            >>> with client.start_run(project="mnist") as run:
            ...     for epoch in range(10):
            ...         for step, batch in enumerate(dataloader):
            ...             loss = train_step(batch)
            ...             run.log({"train/loss": loss, "epoch": epoch, "step": step})
            ...         val_loss, val_acc = evaluate(model)
            ...         run.log({"eval/loss": val_loss, "eval/accuracy": val_acc})
        """
        ipc = self._get_ipc()

        # Auto-transition to "running" on first log if not already running
        if not self._is_running[0]:
            if ipc is not None:
                ipc.send({"type": "status", "status": RunStatus.RUNNING})
            else:
                self._client.update_status(self.id, RunStatus.RUNNING)
            self._is_running[0] = True

        now_ns = time.monotonic_ns()
        relative_time_ns = now_ns - self._start_time_ns

        if ipc is not None:
            ipc.send(
                {
                    "type": "log",
                    "values": values,
                    "timestamp": now_ns,
                    "relative_time_ns": relative_time_ns,
                }
            )
        else:
            self._client.log_metrics(self.id, values, _at=now_ns, _relative_time_ns=relative_time_ns)

    def __enter__(self) -> "Run":
        ipc = self._get_ipc()
        if not self._is_running[0]:
            if ipc is not None:
                ipc.send({"type": "status", "status": RunStatus.RUNNING})
            else:
                self._client.update_status(self.id, RunStatus.RUNNING)
            self._is_running[0] = True

        # Install console capture (stdout/stderr tee + logging handler)
        if self._capture_output and ipc is not None:
            stdout_tee = _TeeWriter(sys.stdout, "console/stdout", ipc, self._start_time_ns)
            stderr_tee = _TeeWriter(sys.stderr, "console/stderr", ipc, self._start_time_ns)
            self._tee_writers[0] = stdout_tee
            self._tee_writers[1] = stderr_tee
            sys.stdout = stdout_tee  # type: ignore[assignment]
            sys.stderr = stderr_tee  # type: ignore[assignment]

            handler = TykoLogHandler(ipc=ipc, start_time_ns=self._start_time_ns)
            self._log_handler[0] = handler
            logging.root.addHandler(handler)

        return self

    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
        # Restore console capture (flush partial lines, then restore originals)
        stdout_tee = self._tee_writers[0]
        stderr_tee = self._tee_writers[1]
        if stdout_tee is not None:
            stdout_tee.flush()
            sys.stdout = stdout_tee._original  # type: ignore[assignment]
            self._tee_writers[0] = None
        if stderr_tee is not None:
            stderr_tee.flush()
            sys.stderr = stderr_tee._original  # type: ignore[assignment]
            self._tee_writers[1] = None
        handler = self._log_handler[0]
        if handler is not None:
            logging.root.removeHandler(handler)
            self._log_handler[0] = None

        # Determine final status: stopped takes priority if sentinel exists
        from ._sidecar import stop_sentinel_path

        sentinel = stop_sentinel_path(client_uuid=self._client_uuid)
        if os.path.exists(sentinel):
            status = RunStatus.STOPPED
            try:
                os.unlink(sentinel)
            except OSError:
                pass
        elif exc_type is not None:
            status = RunStatus.FAILED
        else:
            status = RunStatus.COMPLETED

        at = datetime.now(timezone.utc)
        ipc = self._get_ipc()

        if ipc is not None:
            ok = ipc.send_and_wait({"type": "close", "status": status, "at": at.isoformat()})
            if not ok:
                # Fallback: sidecar unresponsive — close directly via HTTP
                self._client.close_run(self.id, status, at)
            # Ask sidecar to exit
            ipc.send_and_wait({"type": "shutdown"})
            ipc.close()
        else:
            self._client.close_run(self.id, status, at)
