# Copyright (C) 2026 GoodData Corporation

"""Utilities for consistent MCP progress reporting."""

from __future__ import annotations

import asyncio
import time
from contextlib import suppress
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any

from aida_mcp.output.policy import MAX_PROGRESS_MESSAGE_CHARS

if TYPE_CHECKING:
    from mcp.server.fastmcp import Context as MCPContext

    Context = MCPContext[Any, Any, Any]
else:  # pragma: no cover - typing guard
    Context = Any


@dataclass(slots=True)
class ProgressTracker:
    """Helper for combined ctx.report_progress + ctx.info updates."""

    ctx: Context | None
    total_units: float
    message_prefix: str = ""
    heartbeat_interval: float = 10.0
    _completed_units: float = field(default=0.0, init=False)
    _lock: asyncio.Lock = field(default_factory=asyncio.Lock, init=False)
    _last_heartbeat: float = field(default=0.0, init=False)
    _closed: bool = field(default=False, init=False)

    def _has_progress(self) -> bool:
        return bool(self.ctx) and self.total_units > 0 and not self._closed

    def _format_label(self, label: str) -> str:
        prefix = self.message_prefix.strip()
        cleaned = label.strip()
        if prefix and cleaned:
            return _truncate_message(f"{prefix} {cleaned}")
        if cleaned:
            return _truncate_message(cleaned)
        return _truncate_message(prefix or "progress")

    async def advance(self, label: str, increment: float = 1.0) -> None:
        """Advance progress and emit combined update."""

        if not self._has_progress():
            return

        async with self._lock:
            self._completed_units = min(self.total_units, self._completed_units + increment)
            completed = self._completed_units

        # IMPORTANT:
        # Do NOT emit the final progress notification (completed == total_units).
        # Cursor (and some other MCP clients) can invalidate the progress token as soon as the tool
        # result is received, and the "100%" progress update often races with that completion.
        # The tool result itself is the completion signal.
        if completed >= self.total_units:
            await self.complete(label)
            return

        await self._emit(label, completed)

    async def heartbeat(self, label: str, *, force: bool = False) -> None:
        """Send a heartbeat without increasing progress (throttled)."""

        if not self._has_progress():
            return

        now = time.monotonic()
        if not force and now - self._last_heartbeat < self.heartbeat_interval:
            return

        async with self._lock:
            completed = self._completed_units
        await self._emit(label, completed)
        self._last_heartbeat = now

    async def complete(self, label: str = "Completed") -> None:
        """Mark progress as complete and close the tracker.

        Note: This intentionally does NOT send a final progress notification.
        Progress notifications sent right before the tool returns often fail with
        "unknown token" errors because the client invalidates the progress token
        as soon as it receives the tool result. The tool result itself signals completion.
        """
        if not self._has_progress():
            return

        async with self._lock:
            self._completed_units = self.total_units
            self._closed = True
        # Intentionally not emitting final notification to avoid race condition

    async def _emit(self, label: str, completed: float) -> None:
        """Send the actual MCP updates."""

        if not self.ctx:
            return

        message = self._format_label(label)

        # IMPORTANT: keep report_progress and info separate.
        #
        # Some MCP clients do not support progress tokens, or can invalidate the token mid-run.
        # In that case report_progress may raise - we still want to emit ctx.info heartbeats so
        # the client sees forward progress.
        with suppress(Exception):
            await self.ctx.report_progress(
                progress=completed,
                total=self.total_units,
                message=message,
            )

        with suppress(Exception):
            total_display = int(self.total_units)
            completed_display = int(completed)
            await self.ctx.info(f"{message} ({completed_display}/{total_display})")

    @property
    def completed_units(self) -> float:
        return self._completed_units

    @property
    def remaining_units(self) -> float:
        return max(self.total_units - self._completed_units, 0.0)


def _truncate_message(message: str) -> str:
    """Limit progress/info messages to a safe length for MCP clients."""
    max_chars = MAX_PROGRESS_MESSAGE_CHARS
    if max_chars <= 0:
        return ""
    if len(message) <= max_chars:
        return message
    if max_chars <= 3:
        return message[:max_chars]
    return message[: max_chars - 3] + "..."
