"""CI-1T high-level monitoring.

Monitor class: auto-manages fleet sessions, buffers scores, and pushes
rounds when enough data accumulates. Supports callbacks for stability events.

watch() decorator: wraps a function so its return value is automatically
tracked as a prediction score.
"""

from __future__ import annotations

import functools
import logging
import threading
from typing import Any, Callable

from .client import Client
from .types import FleetSnapshot, RoundResult, to_q16

logger = logging.getLogger("ci1t")

# Default episode size (scores per episode per node)
DEFAULT_EPISODE_SIZE = 3


class Monitor:
    """Fleet session monitor with auto-buffering and callbacks.

    Creates a fleet session on first push and auto-manages the lifecycle.
    Buffers scores per node and flushes to the API when a full episode
    accumulates (default: 3 scores per node).

    Usage::

        monitor = ci1t.Monitor("my-fleet", nodes=["gpt-4o", "claude"])

        # Push one score per node at a time
        monitor.push({"gpt-4o": 0.92, "claude": 0.87})
        monitor.push({"gpt-4o": 0.89, "claude": 0.85})
        monitor.push({"gpt-4o": 0.91, "claude": 0.86})
        # ^ Third push triggers a round flush

        # Or push a full episode at once
        monitor.push({
            "gpt-4o": [0.92, 0.89, 0.91],
            "claude": [0.87, 0.85, 0.86],
        })

        state = monitor.state()
        monitor.close()  # deletes the session

    Args:
        name: Session name (for your reference, not sent to API).
        nodes: List of node names. Required.
        api_key: CI-1T API key. Falls back to CI1T_API_KEY env var.
        base_url: API base URL override.
        episode_size: Scores per episode per node (default 3).
        on_drift: Callback fired when any node enters Drifting+ status.
            Signature: (node_name: str, snapshot: FleetSnapshot) -> None
        on_ghost: Callback fired when a ghost is confirmed.
            Signature: (node_name: str, snapshot: FleetSnapshot) -> None
        on_round: Callback fired after every successful round push.
            Signature: (result: RoundResult) -> None
    """

    def __init__(
        self,
        name: str,
        nodes: list[str],
        *,
        api_key: str | None = None,
        base_url: str | None = None,
        episode_size: int = DEFAULT_EPISODE_SIZE,
        on_drift: Callable[[str, FleetSnapshot], None] | None = None,
        on_ghost: Callable[[str, FleetSnapshot], None] | None = None,
        on_round: Callable[[RoundResult], None] | None = None,
    ) -> None:
        if not nodes:
            raise ValueError("At least one node name is required.")
        if episode_size < 3:
            raise ValueError("episode_size must be >= 3.")

        self.name = name
        self.nodes = list(nodes)
        self.episode_size = episode_size

        self._client = Client(api_key=api_key, base_url=base_url)
        self._session_id: str | None = None
        self._buffers: dict[str, list[int]] = {n: [] for n in self.nodes}
        self._lock = threading.Lock()
        self._round_count = 0
        self._closed = False

        # Callbacks
        self._on_drift = on_drift
        self._on_ghost = on_ghost
        self._on_round = on_round

    @property
    def session_id(self) -> str | None:
        """Current fleet session ID, or None if not yet created."""
        return self._session_id

    def push(self, scores: dict[str, float | int | list[float | int]]) -> RoundResult | None:
        """Push scores for one or more nodes.

        Accepts either a single score per node (buffered until episode_size
        scores accumulate) or a list of scores per node. If any node's buffer
        reaches episode_size, a round is flushed to the API.

        Args:
            scores: Dict mapping node names to either a single score or a
                list of scores. Scores can be floats (0.0-1.0) or Q0.16
                ints (0-65535).

        Returns:
            RoundResult if a round was flushed, None if still buffering.

        Raises:
            ValueError: If a node name is not recognized.
            RuntimeError: If the monitor has been closed.
        """
        if self._closed:
            raise RuntimeError("Monitor has been closed.")

        with self._lock:
            # Buffer incoming scores
            for node_name, value in scores.items():
                if node_name not in self._buffers:
                    raise ValueError(
                        f"Unknown node '{node_name}'. "
                        f"Known nodes: {self.nodes}"
                    )
                if isinstance(value, (list, tuple)):
                    self._buffers[node_name].extend(to_q16(s) for s in value)
                else:
                    self._buffers[node_name].append(to_q16(value))

            # Check if all nodes have enough scores for a flush
            min_buffered = min(len(buf) for buf in self._buffers.values())
            if min_buffered < self.episode_size:
                return None

            # Extract one episode per node
            round_scores: list[list[int]] = []
            for node_name in self.nodes:
                buf = self._buffers[node_name]
                episode = buf[: self.episode_size]
                self._buffers[node_name] = buf[self.episode_size :]
                round_scores.append(episode)

            # Ensure session exists
            if self._session_id is None:
                info = self._client.session_create(
                    node_count=len(self.nodes),
                    node_names=self.nodes,
                )
                self._session_id = info.session_id
                logger.info("Created fleet session %s for '%s'", self._session_id, self.name)

            # Push the round
            result = self._client.session_round(self._session_id, round_scores)
            self._round_count += 1

        # Fire callbacks outside the lock
        self._fire_callbacks(result)

        return result

    def flush(self) -> RoundResult | None:
        """Force-flush any buffered scores, padding with the last score if needed.

        Useful for cleanup when you want to push partial buffers.
        Returns None if all buffers are empty.
        """
        if self._closed:
            raise RuntimeError("Monitor has been closed.")

        with self._lock:
            max_buffered = max(len(buf) for buf in self._buffers.values())
            if max_buffered == 0:
                return None

            # Pad each buffer to episode_size using the last value
            round_scores: list[list[int]] = []
            for node_name in self.nodes:
                buf = self._buffers[node_name]
                if not buf:
                    # No data at all for this node - use midpoint
                    buf = [32768]
                while len(buf) < self.episode_size:
                    buf.append(buf[-1])
                episode = buf[: self.episode_size]
                self._buffers[node_name] = buf[self.episode_size :]
                round_scores.append(episode)

            if self._session_id is None:
                info = self._client.session_create(
                    node_count=len(self.nodes),
                    node_names=self.nodes,
                )
                self._session_id = info.session_id

            result = self._client.session_round(self._session_id, round_scores)
            self._round_count += 1

        self._fire_callbacks(result)
        return result

    def state(self) -> FleetSnapshot | None:
        """Get the current fleet session state.

        Returns:
            FleetSnapshot if a session exists, None otherwise.
        """
        if self._session_id is None:
            return None
        return self._client.session_state(self._session_id)

    def close(self) -> None:
        """Flush remaining scores and delete the fleet session."""
        if self._closed:
            return
        self._closed = True

        try:
            self.flush()
        except Exception:
            logger.debug("Flush on close failed", exc_info=True)

        if self._session_id:
            try:
                self._client.session_delete(self._session_id)
                logger.info("Deleted fleet session %s", self._session_id)
            except Exception:
                logger.debug("Session delete on close failed", exc_info=True)

        self._client.close()

    def __enter__(self) -> Monitor:
        return self

    def __exit__(self, *exc: object) -> None:
        self.close()

    def __repr__(self) -> str:
        status = "closed" if self._closed else "active"
        return (
            f"Monitor(name={self.name!r}, nodes={self.nodes}, "
            f"rounds={self._round_count}, status={status})"
        )

    def _fire_callbacks(self, result: RoundResult) -> None:
        """Fire registered callbacks based on round results."""
        snapshot = result.snapshot

        if self._on_round:
            try:
                self._on_round(result)
            except Exception:
                logger.warning("on_round callback failed", exc_info=True)

        for i, node in enumerate(snapshot.nodes):
            name = self.nodes[i] if i < len(self.nodes) else f"node-{i}"

            if self._on_drift and not node.stable:
                try:
                    self._on_drift(name, snapshot)
                except Exception:
                    logger.warning("on_drift callback failed for %s", name, exc_info=True)

            if self._on_ghost and node.ghost_confirmed:
                try:
                    self._on_ghost(name, snapshot)
                except Exception:
                    logger.warning("on_ghost callback failed for %s", name, exc_info=True)


# ── watch() decorator ──


def watch(
    name: str = "default",
    *,
    api_key: str | None = None,
    base_url: str | None = None,
    episode_size: int = DEFAULT_EPISODE_SIZE,
    on_drift: Callable | None = None,
    on_ghost: Callable | None = None,
    on_round: Callable | None = None,
) -> Callable:
    """Decorator that auto-tracks a function's return value as a prediction score.

    The decorated function must return a float (0.0-1.0) or int (0-65535).
    Scores are buffered and pushed to CI-1T as single-node fleet episodes.

    Usage::

        @ci1t.watch("my-classifier")
        def predict(text: str) -> float:
            return model.predict_proba(text)[1]

        # Each call to predict() returns normally AND buffers the score.
        # Every 3 calls, scores are pushed to CI-1T for stability analysis.
        result = predict("hello world")

    Args:
        name: Model/pipeline name for this monitor.
        api_key: CI-1T API key override.
        base_url: API base URL override.
        episode_size: Scores per episode (default 3).
        on_drift: Callback when stability drifts.
        on_ghost: Callback when ghost is confirmed.
        on_round: Callback after each round push.
    """

    def decorator(fn: Callable) -> Callable:
        monitor = Monitor(
            name=name,
            nodes=[name],
            api_key=api_key,
            base_url=base_url,
            episode_size=episode_size,
            on_drift=on_drift,
            on_ghost=on_ghost,
            on_round=on_round,
        )

        @functools.wraps(fn)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            result = fn(*args, **kwargs)

            # Only track numeric returns
            if isinstance(result, (int, float)):
                try:
                    monitor.push({name: result})
                except Exception:
                    logger.warning("ci1t.watch push failed", exc_info=True)

            return result

        # Expose the monitor for manual inspection
        wrapper.monitor = monitor  # type: ignore[attr-defined]
        return wrapper

    return decorator
