"""
Speed profiling tools for Local Deep Research.

This module provides functionality for measuring execution time
of different components and processes in the research system.
"""

from loguru import logger
import time
from contextlib import contextmanager
from typing import Any, Callable, Dict


class SpeedProfiler:
    """
    Profiler for tracking execution speed of components.

    This class provides methods for timing operations and
    collecting performance statistics for later analysis.
    """

    def __init__(self):
        """Initialize the profiler with empty timing data."""
        self.timings = {}
        self.current_timers = {}
        self.total_start_time = None
        self.total_end_time = None

    def start(self):
        """Start the global profiling session."""
        self.timings = {}
        self.current_timers = {}
        self.total_start_time = time.time()

    def stop(self):
        """Stop the global profiling session."""
        self.total_end_time = time.time()

        # Stop any timers that are still running
        for name in list(self.current_timers.keys()):
            self.stop_timer(name)

    def start_timer(self, name: str):
        """
        Start a named timer.

        Args:
            name: Name of the timer to start
        """
        if name in self.current_timers:
            logger.warning(f"Timer '{name}' is already running. Restarting.")

        self.current_timers[name] = time.time()

    def stop_timer(self, name: str):
        """
        Stop a named timer and record the elapsed time.

        Args:
            name: Name of the timer to stop
        """
        if name not in self.current_timers:
            logger.warning(f"Timer '{name}' was not started.")
            return

        elapsed = time.time() - self.current_timers[name]

        if name not in self.timings:
            self.timings[name] = {
                "total": elapsed,
                "count": 1,
                "min": elapsed,
                "max": elapsed,
                "starts": [self.current_timers[name]],
                "durations": [elapsed],
            }
        else:
            self.timings[name]["total"] += elapsed
            self.timings[name]["count"] += 1
            self.timings[name]["min"] = min(self.timings[name]["min"], elapsed)
            self.timings[name]["max"] = max(self.timings[name]["max"], elapsed)
            self.timings[name]["starts"].append(self.current_timers[name])
            self.timings[name]["durations"].append(elapsed)

        del self.current_timers[name]

    @contextmanager
    def timer(self, name: str):
        """
        Context manager for timing a block of code.

        Args:
            name: Name of the timer

        Example:
            with profiler.timer("my_operation"):
                # Code to time
                do_something()
        """
        self.start_timer(name)
        try:
            yield
        finally:
            self.stop_timer(name)

    def get_timings(self) -> Dict[str, Any]:
        """
        Get all recorded timings.

        Returns:
            Dictionary of timing data for all measured operations
        """
        result = self.timings.copy()

        # Add averages
        for data in result.values():
            if data["count"] > 0:
                data["avg"] = data["total"] / data["count"]

        # Add total duration
        if (
            self.total_start_time is not None
            and self.total_end_time is not None
        ):
            result["total"] = {
                "total": self.total_end_time - self.total_start_time,
                "count": 1,
                "min": self.total_end_time - self.total_start_time,
                "max": self.total_end_time - self.total_start_time,
                "avg": self.total_end_time - self.total_start_time,
                "starts": [self.total_start_time],
                "durations": [self.total_end_time - self.total_start_time],
            }

        return result

    def get_summary(self) -> Dict[str, float]:
        """
        Get a summary of timing information.

        Returns:
            Dictionary with summary statistics
        """
        timings = self.get_timings()
        summary = {}

        # Total duration
        if "total" in timings:
            summary["total_duration"] = timings["total"]["total"]
        elif (
            self.total_start_time is not None
            and self.total_end_time is not None
        ):
            summary["total_duration"] = (
                self.total_end_time - self.total_start_time
            )
        else:
            summary["total_duration"] = sum(
                t["total"] for t in timings.values()
            )

        # Component durations
        for name, data in timings.items():
            if name != "total":
                summary[f"{name}_duration"] = data["total"]
                summary[f"{name}_percent"] = (
                    data["total"] / summary["total_duration"] * 100
                    if summary["total_duration"] > 0
                    else 0
                )

        # Per-operation breakdowns
        for name, data in timings.items():
            if data["count"] > 0:
                summary[f"{name}_per_operation"] = data["total"] / data["count"]

        return summary

    def print_summary(self):
        """Print a formatted summary of timing information."""
        summary = self.get_summary()
        total = summary.get("total_duration", 0)

        print("\n===== SPEED PROFILE SUMMARY =====")
        print(f"Total execution time: {total:.2f} seconds")
        print("\n--- Component Breakdown ---")

        # Print each component's timing
        for name, data in self.timings.items():
            if name != "total":
                percent = data["total"] / total * 100 if total > 0 else 0
                print(
                    f"{name}: {data['total']:.2f}s ({percent:.1f}%) - "
                    f"{data['count']} calls, avg {data['total'] / data['count']:.3f}s per call"
                )

        print("\n==============================")


def time_function(func: Callable) -> Callable:
    """
    Decorator to time a function's execution.

    Args:
        func: Function to time

    Returns:
        Wrapped function that logs its execution time

    Example:
        @time_function
        def my_slow_function():
            # Some slow code
            pass
    """

    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        elapsed = time.time() - start_time

        logger.info(f"{func.__name__} took {elapsed:.3f} seconds")

        return result

    return wrapper
