import atexit
import logging
import os
import re
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)


_ANSI_ESCAPE = re.compile(r"\x1b\[[0-9;]*[A-Za-z]|\x1b\][^\x07]*\x07|\x1b[()][AB012]")


def _strip_ansi(text: str) -> str:
    return _ANSI_ESCAPE.sub("", text)


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)
                # tqdm/progress bars use \r to overwrite — keep only the last segment
                if "\r" in line:
                    line = line.rsplit("\r", 1)[-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:
        clean = _strip_ansi(line)
        if not clean:
            return
        now_ns = time.time_ns()
        self._ipc.send(
            {
                "type": "console",
                "stream": self._key,
                "text": clean,
                "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 = _strip_ansi(self.format(record))
            if not line:
                return
            now_ns = time.time_ns()
            self._ipc.send(
                {
                    "type": "console",
                    "stream": "stdout",
                    "text": line,
                    "timestamp": now_ns,
                    "relative_time_ns": time.monotonic_ns() - self._start_time_ns,
                }
            )
        except Exception:
            self.handleError(record)


def _format_duration(ns: int) -> str:
    """Format nanoseconds into a human-readable duration string."""
    total_seconds = ns / 1e9
    if total_seconds < 60:
        return f"{total_seconds:.1f}s"
    minutes = int(total_seconds // 60)
    seconds = int(total_seconds % 60)
    if minutes < 60:
        return f"{minutes}m {seconds}s"
    hours = minutes // 60
    minutes = minutes % 60
    return f"{hours}h {minutes}m {seconds}s"


@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)
    _silent: bool = field(default=False, 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)
    # Guard against double-close (atexit + explicit __exit__)
    _closed: list[bool] = field(init=False, repr=False)
    # Set to True by the sys.excepthook chain when an unhandled exception occurs
    _crashed: list[bool] = field(init=False, repr=False)
    # [original_hook, chained_hook] — stored so __exit__ can restore the original
    _excepthook: list[Any] = field(init=False, repr=False)
    # Counter for metric log calls
    _log_count: list[int] = 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, "_closed", [False])
        object.__setattr__(self, "_crashed", [False])
        object.__setattr__(self, "_excepthook", [None, None])  # [original, chained]
        object.__setattr__(self, "_log_count", [0])
        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]

    def _print(self, msg: str) -> None:
        """Print a tyko: prefixed message to stderr, unless silent."""
        if self._silent:
            return
        # Use the original stderr if we've installed a tee writer
        stream = self._tee_writers[1]._original if self._tee_writers[1] is not None else sys.stderr
        print(f"tyko: {msg}", file=stream, flush=True)

    @property
    def display_name(self) -> str:
        return f"{self.name}-{self.sequential}"

    @property
    def url(self) -> str:
        return f"https://tyko-labs.com/run/{self.project.name}/{self.display_name}"

    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

        self._log_count[0] += 1

        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 log_table(
        self,
        name: str,
        data: object,
        *,
        metadata: dict[str, object] | None = None,
    ) -> str:
        """Upload tabular data as a Parquet artifact.

        Accepts:
        - polars.DataFrame or pandas.DataFrame
        - dict[str, list] (column-oriented)
        - list[dict] (row-oriented records)

        Args:
            name: Collection name for the artifact (e.g. "predictions").
            data: Tabular data in any of the supported formats.
            metadata: Optional key-value metadata.

        Returns:
            The artifact version ID.

        Example:
            >>> run.log_table("predictions", df)
            >>> run.log_table("results", [{"y": 0.5, "pred": 0.6}])
            >>> run.log_table("results", {"y": [0.5], "pred": [0.6]})
        """
        from ._dataframe import upload_table

        return upload_table(
            run=self,
            name=name,
            data=data,
            metadata=metadata,
        )

    def log_series(
        self,
        name: str,
        series: object,
        *,
        metadata: dict[str, object] | None = None,
    ) -> str:
        """Upload a Series as a Parquet artifact.

        Converts the Series to a single-column DataFrame and uploads it.

        Args:
            name: Collection name for the artifact (e.g. "training_loss").
            series: A pandas or polars Series (or any SeriesLike).
            metadata: Optional key-value metadata.

        Returns:
            The artifact version ID.

        Example:
            >>> run.log_series("class_distribution", series)
        """
        from ._dataframe import upload_series

        return upload_series(
            run=self,
            name=name,
            series=series,
            metadata=metadata,
        )

    _CHART_TYPES = frozenset({"line", "scatter", "bar", "histogram", "table"})

    def chart(
        self,
        name: str,
        *,
        table: str,
        type: str,
        x: str,
        y: str | None = None,
        color: str | None = None,
    ) -> None:
        """Define a chart visualization for a logged table.

        The chart definition is sent to the server and rendered by the
        frontend on the run page.  No data is sent — only a lightweight
        JSON spec that references a table by name and column names.

        Args:
            name: Display name for the chart (e.g. "loss_curve").
            table: Name of a table previously logged with log_table().
            type: Chart type — "line", "scatter", "bar", "histogram", or "table".
            x: Column name for the x-axis.
            y: Column name for the y-axis (optional for histogram).
            color: Optional column name used for colour grouping.

        Example:
            >>> run.log_table("training", df)
            >>> run.chart("loss_curve", table="training", type="line", x="epoch", y="loss")
            >>> run.chart("acc_vs_loss", table="training", type="scatter", x="loss", y="accuracy")
        """
        if type not in self._CHART_TYPES:
            raise ValueError(f"Invalid chart type {type!r}. Must be one of: {', '.join(sorted(self._CHART_TYPES))}")

        self._client.create_chart(
            self.id,
            name=name,
            table=table,
            chart_type=type,
            x=x,
            y=y,
            color=color,
        )

    def log_artifact(
        self,
        path: str | os.PathLike[str],
        name: str | None = None,
        artifact_type: str = "model",
        metadata: dict[str, object] | None = None,
    ) -> str:
        """Upload a file or directory as an artifact attached to this run.

        Scans the path (file or directory) and uploads all files to the
        artifact storage backend. Returns the artifact version ID.

        Args:
            path: Path to a file or directory to upload.
            name: Collection name for the artifact. Defaults to the
                  basename of the path.
            artifact_type: Type label (e.g. "model", "checkpoint", "dataset").
            metadata: Optional key-value metadata stored as filter_data
                      on the artifact version (e.g. {"epoch": 5, "wer": 0.12}).

        Returns:
            The artifact version ID.

        Example:
            >>> version_id = run.log_artifact("output/merged", name="whisper-fine-tuned")
            >>> version_id = run.log_artifact("results.csv", artifact_type="evaluation")
        """
        import mimetypes
        from pathlib import Path as _Path

        artifact_path = _Path(path)
        if not artifact_path.exists():
            raise FileNotFoundError(f"Artifact path does not exist: {artifact_path}")

        if name is None:
            name = artifact_path.name

        # Collect files to upload
        if artifact_path.is_file():
            files = [artifact_path]
            base = artifact_path.parent
        else:
            files = sorted(f for f in artifact_path.rglob("*") if f.is_file())
            base = artifact_path

        if not files:
            raise ValueError(f"No files found in {artifact_path}")

        # Create pending version
        version_id = self._client.create_artifact_version(
            self.id,
            name,
            artifact_type,
            filter_data=metadata,
        )

        try:
            for file in files:
                relative = str(file.relative_to(base))
                content_type = mimetypes.guess_type(str(file))[0] or "application/octet-stream"
                data = file.read_bytes()
                self._client.upload_artifact_file(version_id, relative, data, content_type)

            self._client.complete_artifact_version(version_id)
        except Exception:
            self._client.fail_artifact_version(version_id)
            raise

        return version_id

    def _install_console_capture(self) -> None:
        """Install stdout/stderr tee writers and logging handler.

        Called from start_run() so console capture works regardless of whether
        the context manager is used. __enter__ calls this too but it's idempotent.
        """
        ipc = self._get_ipc()
        if not self._capture_output or ipc is None:
            return
        # Already installed — don't double-wrap
        if self._tee_writers[0] is not None:
            return
        stdout_tee = _TeeWriter(sys.stdout, "stdout", ipc, self._start_time_ns)
        stderr_tee = _TeeWriter(sys.stderr, "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)

    def finish(self, *, failed: bool = False) -> None:
        """Close the run and send its final status to the server.

        Call this in Jupyter notebooks when you are done with a run.
        It is safe to call multiple times (idempotent).

        Args:
            failed: If True, mark the run as failed instead of completed.

        Example:
            >>> run = client.start_run(project="mnist")
            >>> # ... training ...
            >>> run.finish()       # mark as completed
            >>> run.finish(failed=True)  # mark as failed
        """
        exc_type = Exception if failed else None
        self.__exit__(exc_type, None, None)

    def __del__(self) -> None:
        """Close the run when the object is garbage collected.

        Handles notebook patterns like `del run` or cell re-execution that
        reassigns the `run` variable — CPython calls this immediately when
        the reference count drops to zero.
        """
        try:
            if not self._closed[0]:
                self.__exit__(None, None, None)
        except Exception:
            pass

    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

        self._install_console_capture()

        return self

    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
        # Guard: atexit may call __exit__ after the context manager already did
        if self._closed[0]:
            return
        self._closed[0] = True
        # Remove the atexit handler now that we're closing explicitly
        atexit.unregister(self.__exit__)

        # Restore sys.excepthook (only if our chained hook is still installed)
        chained = self._excepthook[1]
        if chained is not None and sys.excepthook is chained:
            sys.excepthook = self._excepthook[0]
        self._excepthook[0] = None
        self._excepthook[1] = 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 or self._crashed[0]:
            status = RunStatus.FAILED
        else:
            status = RunStatus.COMPLETED

        # Print footer summary
        elapsed = _format_duration(time.monotonic_ns() - self._start_time_ns)
        count = self._log_count[0]
        entries = f"{count} metric entr{'y' if count == 1 else 'ies'} logged"
        self._print(f"Run {self.display_name} {status.value} in {elapsed} ({entries})")
        self._print(f"View run at {self.url}")

        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)
