from __future__ import annotations

import asyncio
import contextlib
import logging
import time
from collections.abc import AsyncIterator
from typing import Any

logger = logging.getLogger(__name__)


class AgentWatchdog:
    """Monitors runner progress via heartbeats. Force-closes stuck connections on stall.

    The key property: calling ``force_unstick`` closes the registered httpx
    client from the *outside*, destroying the TCP socket.  The stuck ``await``
    inside httpx receives an immediate, un-ignorable connection error --
    no reliance on ``CancelledError`` propagation.

    Use ``suspend()`` as a context manager around known long-running
    operations (script execution, dataset downloads) where no heartbeats
    are expected::

        async with watchdog.suspend("running_training_script"):
            result = await run_script(...)
    """

    def __init__(self, stall_timeout: float = 150, max_script_warn: float = 600) -> None:
        self._stall_timeout = stall_timeout
        self._max_script_warn = max_script_warn
        self._last_heartbeat: float = time.monotonic()
        self._current_op: str = ""
        self._stall_count: int = 0
        self._client_ref: Any | None = None  # anthropic.AsyncAnthropic
        self._suspended: bool = False
        self._suspend_start: float | None = None
        self._script_warned: bool = False

    # -- runner / coder call these -------------------------------------------

    def heartbeat(self, operation: str = "") -> None:
        """Signal that the agent is making progress."""
        self._last_heartbeat = time.monotonic()
        self._current_op = operation
        logger.debug("Watchdog: heartbeat '%s'", operation)
        if self._stall_count > 0:
            logger.info("Watchdog: heartbeat received after %d stall(s), resetting", self._stall_count)
            self._stall_count = 0

    def register_client(self, client: Any) -> None:
        """Store a reference to the active AsyncAnthropic client."""
        self._client_ref = client

    def clear_client(self) -> None:
        self._client_ref = None

    @contextlib.asynccontextmanager
    async def suspend(self, reason: str = "") -> AsyncIterator[None]:
        """Temporarily disable stall detection for known long-running ops."""
        self._suspended = True
        self._suspend_start = time.monotonic()
        self._script_warned = False
        self._current_op = reason or self._current_op
        logger.debug("Watchdog: suspended (%s)", reason)
        try:
            yield
        finally:
            elapsed = time.monotonic() - (self._suspend_start or time.monotonic())
            self._suspended = False
            self._suspend_start = None
            self._script_warned = False
            self._last_heartbeat = time.monotonic()
            logger.debug("Watchdog: resumed after %s (%.1fs elapsed)", reason, elapsed)

    # -- state queries -------------------------------------------------------

    @property
    def seconds_since_heartbeat(self) -> float:
        return time.monotonic() - self._last_heartbeat

    @property
    def is_stalled(self) -> bool:
        if self._suspended:
            return False
        return self.seconds_since_heartbeat > self._stall_timeout

    # -- intervention --------------------------------------------------------

    async def force_unstick(self) -> None:
        """Force-close the registered client, destroying the TCP socket."""
        self._stall_count += 1
        client = self._client_ref
        if client is None:
            logger.warning(
                "Watchdog: stall detected (%.1fs on '%s', #%d) but no client registered",
                self.seconds_since_heartbeat, self._current_op, self._stall_count,
            )
            return
        try:
            is_closed = client.is_closed() if callable(client.is_closed) else client.is_closed
        except Exception:
            is_closed = False
        if is_closed:
            logger.warning(
                "Watchdog: stall detected (%.1fs on '%s', #%d) but client already closed",
                self.seconds_since_heartbeat, self._current_op, self._stall_count,
            )
            return
        logger.warning(
            "Watchdog: force-closing client after %.1fs stall on '%s' (stall #%d)",
            self.seconds_since_heartbeat, self._current_op, self._stall_count,
        )
        try:
            await client.close()
        except Exception:
            logger.debug("Watchdog: error during force-close (ignored)", exc_info=True)

    # -- supervisor loop -----------------------------------------------------

    async def supervise(self, runner_task: asyncio.Task) -> None:
        """Run alongside the runner task. Checks every 5s; intervenes on stall."""
        logger.info("Watchdog: supervisor started (stall_timeout=%.0fs)", self._stall_timeout)
        check_count = 0
        try:
            while not runner_task.done():
                await asyncio.sleep(5)
                check_count += 1
                elapsed = self.seconds_since_heartbeat
                if check_count % 6 == 0:  # every ~30s
                    logger.debug(
                        "Watchdog: check #%d — op='%s' since_hb=%.0fs suspended=%s client=%s",
                        check_count, self._current_op, elapsed,
                        self._suspended, "yes" if self._client_ref else "no",
                    )
                if self._suspended and self._suspend_start is not None:
                    suspend_elapsed = time.monotonic() - self._suspend_start
                    if suspend_elapsed > self._max_script_warn and not self._script_warned:
                        self._script_warned = True
                        logger.warning(
                            "Watchdog: suspended operation '%s' running for %.0fs (exceeds %.0fs threshold)",
                            self._current_op, suspend_elapsed, self._max_script_warn,
                        )

                if self.is_stalled:
                    await self.force_unstick()
        except asyncio.CancelledError:
            pass
        logger.info("Watchdog: supervisor stopped (runner done=%s)", runner_task.done())
