"""
Profiler - Performance measurement and reporting.

Tracks timing for each phase of inference:
- Disk I/O (layer loading)
- CPU→GPU transfer
- GPU computation (attention, MLP, routing)
- Token generation throughput
"""

from __future__ import annotations

import logging
import time
from collections import defaultdict
from contextlib import contextmanager
from dataclasses import dataclass, field

import torch

logger = logging.getLogger(__name__)


@dataclass
class TimingRecord:
    """Accumulated timing for a named operation."""

    name: str
    total_seconds: float = 0.0
    count: int = 0
    min_seconds: float = float("inf")
    max_seconds: float = 0.0

    @property
    def avg_seconds(self) -> float:
        return self.total_seconds / max(self.count, 1)

    def update(self, elapsed: float) -> None:
        self.total_seconds += elapsed
        self.count += 1
        self.min_seconds = min(self.min_seconds, elapsed)
        self.max_seconds = max(self.max_seconds, elapsed)


class Profiler:
    """
    Hierarchical profiler for inference performance analysis.

    Usage:
        profiler = Profiler(enabled=True)

        with profiler.track("layer_0.attention"):
            # ... attention computation ...

        with profiler.track("layer_0.io"):
            # ... disk loading ...

        print(profiler.report())
    """

    def __init__(self, enabled: bool = False):
        self.enabled = enabled
        self._timings: dict[str, TimingRecord] = defaultdict(
            lambda: TimingRecord(name="")
        )
        self._start_time: float = 0.0
        self._tokens_generated: int = 0
        self._bytes_loaded: int = 0

    def start(self) -> None:
        """Mark the start of generation."""
        self._start_time = time.perf_counter()

    @contextmanager
    def track(self, name: str):
        """Context manager to time a named operation."""
        if not self.enabled:
            yield
            return

        if torch.cuda.is_available():
            torch.cuda.synchronize()
        start = time.perf_counter()

        yield

        if torch.cuda.is_available():
            torch.cuda.synchronize()
        elapsed = time.perf_counter() - start

        record = self._timings[name]
        record.name = name
        record.update(elapsed)

    def record_token(self) -> None:
        """Record that a token was generated."""
        self._tokens_generated += 1

    def record_bytes_loaded(self, nbytes: int) -> None:
        """Record bytes loaded from disk."""
        self._bytes_loaded += nbytes

    def report(self) -> str:
        """Generate a human-readable profiling report."""
        if not self.enabled:
            return "Profiling disabled"

        total_time = time.perf_counter() - self._start_time
        tokens_per_sec = self._tokens_generated / max(total_time, 1e-9)

        lines = [
            "=" * 60,
            "DOKODEMO AI - Performance Report",
            "=" * 60,
            f"Total time: {total_time:.2f}s",
            f"Tokens generated: {self._tokens_generated}",
            f"Throughput: {tokens_per_sec:.2f} tokens/sec",
            f"Data loaded: {self._bytes_loaded / 1e9:.2f} GB",
            "",
            f"{'Operation':<35} {'Total':>8} {'Count':>6} {'Avg':>8} {'%':>6}",
            "-" * 65,
        ]

        # Group by category
        categories: dict[str, float] = defaultdict(float)
        sorted_timings = sorted(
            self._timings.values(),
            key=lambda t: t.total_seconds,
            reverse=True,
        )

        for record in sorted_timings:
            pct = (record.total_seconds / max(total_time, 1e-9)) * 100
            lines.append(
                f"{record.name:<35} "
                f"{record.total_seconds:>7.3f}s "
                f"{record.count:>6d} "
                f"{record.avg_seconds * 1000:>6.1f}ms "
                f"{pct:>5.1f}%"
            )

            # Categorize
            if "io" in record.name or "load" in record.name:
                categories["I/O"] += record.total_seconds
            elif "compute" in record.name or "forward" in record.name:
                categories["Compute"] += record.total_seconds
            elif "transfer" in record.name:
                categories["Transfer"] += record.total_seconds
            else:
                categories["Other"] += record.total_seconds

        lines.extend(["", "Category Summary:"])
        for cat, total in sorted(
            categories.items(), key=lambda x: x[1], reverse=True
        ):
            pct = (total / max(total_time, 1e-9)) * 100
            lines.append(f"  {cat:<15} {total:>7.3f}s  ({pct:.1f}%)")

        lines.append("=" * 60)
        return "\n".join(lines)

    def reset(self) -> None:
        """Reset all profiling data."""
        self._timings.clear()
        self._start_time = 0.0
        self._tokens_generated = 0
        self._bytes_loaded = 0
