# Copyright (C) 2026 GoodData Corporation

"""Timing utilities for validation operations."""

from __future__ import annotations

import time
from contextlib import contextmanager, suppress
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from collections.abc import Iterator

Context = Any


class StepTimer:
    """Tracks timing for multiple validation steps."""

    def __init__(self) -> None:
        self.start_time: float = time.time()
        self.step_timings: list[tuple[str, float]] = []

    def record_step(self, step_name: str, duration: float) -> None:
        """Record a completed step with its duration."""
        self.step_timings.append((step_name, duration))

    def total_duration(self) -> float:
        """Get total elapsed time since timer was created."""
        return time.time() - self.start_time

    def format_summary(self) -> list[str]:
        """Format timing summary for output.

        Returns a prominently formatted summary that MCP clients can easily identify
        and display. The summary is always placed at the end of tool responses.
        """
        if not self.step_timings:
            return []

        lines = [
            "",
            "=" * 60,
            "⏱️  TIMING SUMMARY",
            "=" * 60,
        ]
        for step_name, step_duration in self.step_timings:
            lines.append(f"  {step_name}: {step_duration:.1f}s")
        lines.append("")
        lines.append(f"  Total: {self.total_duration():.1f}s")
        return lines


@contextmanager
def time_step(step_name: str, timer: StepTimer | None = None) -> Iterator[float]:
    """Context manager to time a single step.

    Args:
        step_name: Name of the step being timed
        timer: Optional StepTimer to record the step in

    Yields:
        Start time of the step

    Example:
        timer = StepTimer()
        with time_step("Lint", timer) as start:
            result = await run_lint()
        # Step is automatically recorded in timer
    """
    start_time = time.time()
    try:
        yield start_time
    finally:
        duration = time.time() - start_time
        if timer:
            timer.record_step(step_name, duration)


def format_step_result(result: str, step_name: str, duration: float) -> str:
    """Add timing information to a step result string.

    Args:
        result: Original result string
        step_name: Name of the step
        duration: Duration in seconds

    Returns:
        Result string with timing prepended
    """
    timing_prefix = f"⏱️  {step_name} ({duration:.1f}s)\n"
    # Try to insert before "Command:" if present, otherwise prepend
    if "Command: " in result:
        return result.replace("Command: ", f"{timing_prefix}Command: ", 1)
    return f"{timing_prefix}{result}"


async def emit_timing_summary(ctx: Context | None, timer: StepTimer) -> None:
    """Send timing summary via ctx.info for clients that display streaming logs."""

    if not ctx:
        return

    summary = timer.format_summary()
    if not summary:
        return

    for line in summary:
        if not line.strip():
            continue
        with suppress(Exception):
            await ctx.info(line)
