"""Lightweight latency profiling for RAG pipelines.

Usage:
    from animuz_core.profiling import Profiler

    p = Profiler()
    p.mark("start")
    # ... do work ...
    p.mark("first_chunk")
    # ... more work ...
    p.mark("tool_call_start")
    p.mark("tool_call_end")
    p.mark("done")
    p.report()

Output:
    === Latency Profile ===
    start → first_chunk        1.234s
    first_chunk → tool_call_start  0.012s
    tool_call_start → tool_call_end  2.456s
    tool_call_end → done       0.089s
    ─────────────────────────
    TOTAL                      3.791s
"""
import time
import json
import logging
from typing import List, Tuple, Optional

logger = logging.getLogger(__name__)


class Profiler:
    """Simple mark-based latency profiler with counters and summaries."""

    def __init__(self, name: str = ""):
        self.name = name
        self.marks: List[Tuple[str, float]] = []
        self._start = time.perf_counter()
        self.counters: dict = {}  # e.g. {"tool_calls": 3}
        self.durations: dict = {}  # e.g. {"mcp_total": [0.5, 0.3]}

    def mark(self, label: str):
        """Record a timing mark."""
        self.marks.append((label, time.perf_counter()))

    def count(self, key: str, increment: int = 1):
        """Increment a named counter."""
        self.counters[key] = self.counters.get(key, 0) + increment

    def record_duration(self, key: str, duration: float):
        """Record a named duration for aggregation."""
        self.durations.setdefault(key, []).append(duration)

    def elapsed(self, label: str) -> Optional[float]:
        """Get elapsed time from start to a specific mark."""
        for name, t in self.marks:
            if name == label:
                return t - self._start
        return None

    def between(self, from_label: str, to_label: str) -> Optional[float]:
        """Get elapsed time between two marks."""
        t_from = t_to = None
        for name, t in self.marks:
            if name == from_label:
                t_from = t
            if name == to_label:
                t_to = t
        if t_from is not None and t_to is not None:
            return t_to - t_from
        return None

    def report(self) -> str:
        """Print and return a formatted latency report."""
        if len(self.marks) < 2:
            return "Not enough marks to report."

        lines = []
        header = f"=== Latency Profile{' (' + self.name + ')' if self.name else ''} ==="
        lines.append(header)

        max_label_len = max(
            len(f"{self.marks[i][0]} → {self.marks[i+1][0]}")
            for i in range(len(self.marks) - 1)
        )

        for i in range(len(self.marks) - 1):
            label = f"{self.marks[i][0]} → {self.marks[i+1][0]}"
            dt = self.marks[i+1][1] - self.marks[i][1]
            lines.append(f"  {label:<{max_label_len}}  {dt:.3f}s")

        total = self.marks[-1][1] - self.marks[0][1]
        lines.append("  " + "─" * (max_label_len + 8))
        lines.append(f"  {'TOTAL':<{max_label_len}}  {total:.3f}s")

        if self.counters:
            lines.append("")
            lines.append("  Counters:")
            for k, v in self.counters.items():
                lines.append(f"    {k}: {v}")

        if self.durations:
            lines.append("")
            lines.append("  Aggregated Durations:")
            for k, vals in self.durations.items():
                total_d = sum(vals)
                avg_d = total_d / len(vals)
                lines.append(f"    {k}: {total_d:.3f}s total, {avg_d:.3f}s avg ({len(vals)} calls)")

        report_str = "\n".join(lines)
        print(report_str)
        return report_str

    def to_dict(self) -> dict:
        """Export marks as a JSON-serializable dict."""
        result = {}
        for i in range(len(self.marks) - 1):
            key = f"{self.marks[i][0]}_to_{self.marks[i+1][0]}"
            result[key] = round(self.marks[i+1][1] - self.marks[i][1], 4)
        if len(self.marks) >= 2:
            result["total"] = round(self.marks[-1][1] - self.marks[0][1], 4)
        if self.counters:
            result["counters"] = dict(self.counters)
        if self.durations:
            result["durations"] = {
                k: {"total": round(sum(v), 4), "avg": round(sum(v)/len(v), 4), "count": len(v)}
                for k, v in self.durations.items()
            }
        return result

    def log(self):
        """Log the profile as structured JSON."""
        logger.info(json.dumps({
            "type": "latency_profile",
            "name": self.name,
            **self.to_dict()
        }))
