"""
Metrics aggregator for releaseops analytics.

Computes aggregate behavioral metrics from raw trace data.
"""

from __future__ import annotations

import logging
import statistics
from typing import Any, Dict, List, Tuple

from llmhq_releaseops.analytics.models import BehaviorMetrics

logger = logging.getLogger(__name__)


class MetricsAggregator:
    """Aggregates raw trace data into BehaviorMetrics."""

    def aggregate(
        self,
        traces: List[Dict[str, Any]],
        bundle_id: str,
        version: str,
        env: str,
    ) -> BehaviorMetrics:
        """
        Compute aggregate metrics from traces.

        Gracefully handles malformed traces by skipping them.
        """
        if not traces:
            return BehaviorMetrics(
                bundle_id=bundle_id,
                bundle_version=version,
                environment=env,
                sample_size=0,
            )

        latencies = []
        token_counts = []
        tool_counts: Dict[str, int] = {}
        errors = 0
        successes = 0
        time_range = self._extract_time_range(traces)

        for trace in traces:
            try:
                lat = self._extract_latency(trace)
                if lat is not None:
                    latencies.append(lat)

                tokens = self._extract_tokens(trace)
                if tokens is not None:
                    token_counts.append(tokens)

                self._extract_tools(trace, tool_counts)

                if self._extract_error(trace):
                    errors += 1
                else:
                    successes += 1
            except Exception as e:
                logger.warning(f"Skipping malformed trace: {e}")

        sample_size = len(traces)
        total_valid = errors + successes
        error_rate = errors / total_valid if total_valid > 0 else 0.0
        success_rate = successes / total_valid if total_valid > 0 else 1.0

        total_tokens = sum(token_counts)
        avg_tokens = total_tokens / len(token_counts) if token_counts else 0.0

        avg_latency = statistics.mean(latencies) if latencies else 0.0
        p50 = self._compute_percentile(latencies, 0.5)
        p95 = self._compute_percentile(latencies, 0.95)
        p99 = self._compute_percentile(latencies, 0.99)

        total_tool_calls = sum(tool_counts.values())
        tool_distribution = {
            name: count / total_tool_calls
            for name, count in tool_counts.items()
        } if total_tool_calls > 0 else {}

        return BehaviorMetrics(
            bundle_id=bundle_id,
            bundle_version=version,
            environment=env,
            sample_size=sample_size,
            time_range=time_range,
            error_rate=error_rate,
            avg_latency_ms=avg_latency,
            p50_latency_ms=p50,
            p95_latency_ms=p95,
            p99_latency_ms=p99,
            total_token_usage=total_tokens,
            avg_token_usage=avg_tokens,
            tool_usage=tool_counts,
            tool_distribution=tool_distribution,
            success_rate=success_rate,
        )

    def _extract_latency(self, trace: Dict[str, Any]) -> float | None:
        """Extract latency in ms from trace."""
        lat = trace.get("latency_ms") or trace.get("duration_ms")
        if lat is not None:
            return float(lat)
        return None

    def _extract_tokens(self, trace: Dict[str, Any]) -> int | None:
        """Extract total token count from trace."""
        total = trace.get("total_tokens")
        if total is not None:
            return int(total)
        input_tokens = trace.get("input_tokens", 0)
        output_tokens = trace.get("output_tokens", 0)
        if input_tokens or output_tokens:
            return int(input_tokens) + int(output_tokens)
        return None

    def _extract_tools(self, trace: Dict[str, Any], tool_counts: Dict[str, int]) -> None:
        """Extract tool usage from trace."""
        tool_calls = trace.get("tool_calls", [])
        if isinstance(tool_calls, list):
            for tc in tool_calls:
                if isinstance(tc, dict):
                    name = tc.get("name", tc.get("function", "unknown"))
                    tool_counts[name] = tool_counts.get(name, 0) + 1

    def _extract_error(self, trace: Dict[str, Any]) -> bool:
        """Check if trace represents an error."""
        if trace.get("error"):
            return True
        status = trace.get("status", "")
        return status in ("error", "failed")

    def _extract_time_range(self, traces: List[Dict[str, Any]]) -> Tuple[str, str]:
        """Determine time range from traces."""
        timestamps = []
        for trace in traces:
            ts = trace.get("timestamp") or trace.get("created_at") or trace.get("start_time")
            if ts:
                timestamps.append(str(ts))
        if not timestamps:
            return ("", "")
        timestamps.sort()
        return (timestamps[0], timestamps[-1])

    def _compute_percentile(self, values: List[float], percentile: float) -> float:
        """Compute percentile value from a list."""
        if not values:
            return 0.0
        if len(values) == 1:
            return values[0]
        sorted_vals = sorted(values)
        idx = percentile * (len(sorted_vals) - 1)
        lower = int(idx)
        upper = min(lower + 1, len(sorted_vals) - 1)
        fraction = idx - lower
        return sorted_vals[lower] + fraction * (sorted_vals[upper] - sorted_vals[lower])
