from dataclasses import asdict, dataclass, field
from os import getenv
from pathlib import Path
from typing import TYPE_CHECKING, Callable, List, Optional
from uuid import uuid4

from memory_profiler import memory_usage

from agno.utils.log import logger, set_log_level_to_debug, set_log_level_to_info
from agno.utils.timer import Timer

if TYPE_CHECKING:
    from rich.console import Console


@dataclass
class PerfResult:
    # Run time performance in seconds
    run_times: List[float] = field(default_factory=list)
    avg_run_time: float = field(init=False)
    min_run_time: float = field(init=False)
    max_run_time: float = field(init=False)
    std_dev_run_time: float = field(init=False)

    # Memory performance in MiB
    memory_usages: List[float] = field(default_factory=list)
    avg_memory_usage: float = field(init=False)
    min_memory_usage: float = field(init=False)
    max_memory_usage: float = field(init=False)
    std_dev_memory_usage: float = field(init=False)

    def __post_init__(self):
        self.compute_stats()

    def compute_stats(self):
        import statistics

        if self.run_times:
            self.avg_run_time = statistics.mean(self.run_times)
            self.min_run_time = min(self.run_times)
            self.max_run_time = max(self.run_times)
            self.std_dev_run_time = statistics.stdev(self.run_times) if len(self.run_times) > 1 else 0

        if self.memory_usages:
            self.avg_memory_usage = statistics.mean(self.memory_usages)
            self.min_memory_usage = min(self.memory_usages)
            self.max_memory_usage = max(self.memory_usages)
            self.std_dev_memory_usage = statistics.stdev(self.memory_usages) if len(self.memory_usages) > 1 else 0

    def print_summary(self, console: Optional["Console"] = None):
        from rich.console import Console
        from rich.table import Table

        if console is None:
            console = Console()

        # Create performance table
        perf_table = Table(title="Performance Summary", show_header=True, header_style="bold magenta")
        perf_table.add_column("Metric", style="cyan")
        perf_table.add_column("Time (seconds)", style="green")
        perf_table.add_column("Memory (MiB)", style="yellow")

        # Add rows
        perf_table.add_row("Average", f"{self.avg_run_time:.6f}", f"{self.avg_memory_usage:.2f}")
        perf_table.add_row("Minimum", f"{self.min_run_time:.6f}", f"{self.min_memory_usage:.2f}")
        perf_table.add_row("Maximum", f"{self.max_run_time:.6f}", f"{self.max_memory_usage:.2f}")
        perf_table.add_row("Std Dev", f"{self.std_dev_run_time:.6f}", f"{self.std_dev_memory_usage:.2f}")

        console.print(perf_table)

    def print_results(self, console: Optional["Console"] = None):
        from rich.console import Console
        from rich.table import Table

        if console is None:
            console = Console()

        # Create runs table
        results_table = Table(title="Individual Runs", show_header=True, header_style="bold magenta")
        results_table.add_column("Run #", style="cyan")
        results_table.add_column("Time (seconds)", style="green")
        results_table.add_column("Memory (MiB)", style="yellow")

        # Add rows
        for i in range(len(self.run_times)):
            results_table.add_row(str(i + 1), f"{self.run_times[i]:.6f}", f"{self.memory_usages[i]:.2f}")

        console.print(results_table)


@dataclass
class PerfEval:
    """Evaluate the performance of a function."""

    # Function to evaluate
    func: Callable

    # Evaluation name
    name: Optional[str] = None
    # Evaluation UUID (autogenerated if not set)
    eval_id: Optional[str] = None

    # Number of iterations to run
    num_iterations: int = 3
    # Result of the evaluation
    result: Optional[PerfResult] = None

    # Print summary of results
    print_summary: bool = False
    # Print detailed results
    print_results: bool = False
    # Save the result to a file
    save_result_to_file: Optional[str] = None

    # debug_mode=True enables debug logs
    debug_mode: bool = False

    def set_eval_id(self) -> str:
        if self.eval_id is None:
            self.eval_id = str(uuid4())
        logger.debug(f"*********** Evaluation ID: {self.eval_id} ***********")
        return self.eval_id

    def set_debug_mode(self) -> None:
        if self.debug_mode or getenv("AGNO_DEBUG", "false").lower() == "true":
            self.debug_mode = True
            set_log_level_to_debug()
            logger.debug("Debug logs enabled")
        else:
            set_log_level_to_info()

    def run(self, *, print_summary: bool = False, print_results: bool = False) -> PerfResult:
        from rich.console import Console
        from rich.live import Live
        from rich.status import Status

        run_times = []
        memory_usages = []
        self.set_eval_id()
        self.set_debug_mode()
        self.print_results = print_results
        self.print_summary = print_summary

        # Add a spinner while running the evaluations
        console = Console()
        with Live(console=console, transient=True) as live_log:
            # Measure run time
            for i in range(self.num_iterations):
                status = Status(
                    f"Running run time evaluation {i + 1}...", spinner="dots", speed=1.0, refresh_per_second=10
                )
                live_log.update(status)

                timer = Timer()
                # Start the timer
                timer.start()
                # Run the function
                self.func()
                # Stop the timer
                timer.stop()
                # Append the time taken
                run_times.append(timer.elapsed)
                logger.debug(f"\nRun {i + 1}:")
                logger.debug(f"Time taken: {timer.elapsed:.6f} seconds")
                status.stop()

            # Measure memory usage
            for i in range(self.num_iterations):
                status = Status(
                    f"Running memory usage evaluation {i + 1}...", spinner="dots", speed=1.0, refresh_per_second=10
                )
                live_log.update(status)
                mem_usage = memory_usage((self.func, (), {}), interval=0.1, max_iterations=1, max_usage=True)
                memory_usages.append(mem_usage)
                logger.debug(f"\nRun {i + 1}:")
                logger.debug(f"Memory usage: {mem_usage} MiB")
                status.stop()

        # Create and store the result
        self.result = PerfResult(run_times=run_times, memory_usages=memory_usages)

        # -*- Save result to file if save_result_to_file is set
        if self.save_result_to_file is not None and self.result is not None:
            try:
                import json

                fn_path = Path(self.save_result_to_file.format(name=self.name, eval_id=self.eval_id))
                if not fn_path.parent.exists():
                    fn_path.parent.mkdir(parents=True, exist_ok=True)
                fn_path.write_text(json.dumps(asdict(self.result), indent=4))
            except Exception as e:
                logger.warning(f"Failed to save result to file: {e}")

        # Show results
        if self.print_results:
            self.result.print_results(console)
        if self.print_summary or self.print_results:
            self.result.print_summary(console)

        logger.debug(f"*********** Evaluation End: {self.eval_id} ***********")
        return self.result
