src/utils/metrics.py

import json
import statistics
import time
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from typing import Any, Dict, List, Optional


@dataclass
class QueryRecord:
    """Record of a single query execution."""
    query: str
    mode: str
    response_time: float
    success: bool
    timestamp: datetime = field(default_factory=datetime.utcnow)
    error: Optional[str] = None
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert record to dictionary."""
        return {
            "query": self.query,
            "mode": self.mode,
            "response_time": self.response_time,
            "success": self.success,
            "timestamp": self.timestamp.isoformat(),
            "error": self.error,
        }


class QueryMetrics:
    """Tracks query execution metrics and statistics."""
    
    def __init__(self, max_history: int = 1000):
        self._max_history = max_history
        self._records: List[QueryRecord] = []
        self._rewrite_records: List[Dict[str, Any]] = []
        self._summarization_records: List[Dict[str, Any]] = []
        self._filtering_records: List[Dict[str, Any]] = []
        self._context_errors: List[Dict[str, Any]] = []
        self._start_time = datetime.utcnow()
    
    @property
    def records(self) -> List[QueryRecord]:
        """Get all query records."""
        return self._records.copy()
    
    @property
    def total_queries(self) -> int:
        """Get total number of queries tracked."""
        return len(self._records)
    
    @property
    def successful_queries(self) -> int:
        """Get number of successful queries."""
        return sum(1 for r in self._records if r.success)
    
    @property
    def failed_queries(self) -> int:
        """Get number of failed queries."""
        return sum(1 for r in self._records if not r.success)
    
    @property
    def success_rate(self) -> float:
        """Get success rate as percentage."""
        if not self._records:
            return 100.0
        return (self.successful_queries / self.total_queries) * 100
    
    @property
    def uptime_seconds(self) -> float:
        """Get seconds since metrics started tracking."""
        return (datetime.utcnow() - self._start_time).total_seconds()
    
    def record(
        self,
        query: str,
        mode: str,
        response_time: float,
        success: bool,
        error: Optional[str] = None,
    ) -> QueryRecord:
        """Record a query execution."""
        record = QueryRecord(
            query=query,
            mode=mode,
            response_time=response_time,
            success=success,
            error=error,
        )
        
        self._records.append(record)
        
        if len(self._records) > self._max_history:
            self._records = self._records[-self._max_history:]
        
        return record
    
    def get_average_response_time(self, mode: Optional[str] = None) -> float:
        """Get average response time in seconds."""
        filtered = self._filter_by_mode(mode)
        if not filtered:
            return 0.0
        return statistics.mean(r.response_time for r in filtered)
    
    def get_p95_latency(self, mode: Optional[str] = None) -> float:
        """Get 95th percentile latency in seconds."""
        filtered = self._filter_by_mode(mode)
        if not filtered:
            return 0.0
        times = sorted(r.response_time for r in filtered)
        index = int(len(times) * 0.95)
        return times[min(index, len(times) - 1)]
    
    def get_p99_latency(self, mode: Optional[str] = None) -> float:
        """Get 99th percentile latency in seconds."""
        filtered = self._filter_by_mode(mode)
        if not filtered:
            return 0.0
        times = sorted(r.response_time for r in filtered)
        index = int(len(times) * 0.99)
        return times[min(index, len(times) - 1)]
    
    def get_min_response_time(self, mode: Optional[str] = None) -> float:
        """Get minimum response time in seconds."""
        filtered = self._filter_by_mode(mode)
        if not filtered:
            return 0.0
        return min(r.response_time for r in filtered)
    
    def get_max_response_time(self, mode: Optional[str] = None) -> float:
        """Get maximum response time in seconds."""
        filtered = self._filter_by_mode(mode)
        if not filtered:
            return 0.0
        return max(r.response_time for r in filtered)
    
    def get_query_count_by_mode(self) -> Dict[str, int]:
        """Get query count grouped by mode."""
        counts: Dict[str, int] = {}
        for record in self._records:
            counts[record.mode] = counts.get(record.mode, 0) + 1
        return counts
    
    def get_error_count_by_type(self) -> Dict[str, int]:
        """Get error count grouped by error message."""
        counts: Dict[str, int] = {}
        for record in self._records:
            if record.error:
                counts[record.error] = counts.get(record.error, 0) + 1
        return counts
    
    def get_queries_per_minute(self) -> float:
        """Get average queries per minute since start."""
        uptime_minutes = self.uptime_seconds / 60
        if uptime_minutes < 0.001:
            return 0.0
        return self.total_queries / uptime_minutes
    
    def get_summary(self) -> Dict[str, Any]:
        """Get complete metrics summary."""
        return {
            "total_queries": self.total_queries,
            "successful_queries": self.successful_queries,
            "failed_queries": self.failed_queries,
            "success_rate": round(self.success_rate, 2),
            "average_response_time": round(self.get_average_response_time(), 3),
            "p95_latency": round(self.get_p95_latency(), 3),
            "p99_latency": round(self.get_p99_latency(), 3),
            "min_response_time": round(self.get_min_response_time(), 3),
            "max_response_time": round(self.get_max_response_time(), 3),
            "queries_per_minute": round(self.get_queries_per_minute(), 2),
            "query_count_by_mode": self.get_query_count_by_mode(),
            "uptime_seconds": round(self.uptime_seconds, 0),
        }
    
    def record_query_rewrite(self, rewritten: bool, method: str, time_ms: float, cache_hit: bool) -> None:
        self._rewrite_records.append({
            "rewritten": rewritten,
            "method": method,
            "time_ms": time_ms,
            "cache_hit": cache_hit,
            "timestamp": datetime.utcnow(),
        })
        self._prune_context_records(self._rewrite_records)

    def record_conversation_summary(self, message_count_before: int, message_count_after: int, time_ms: float) -> None:
        ratio = 1 - (message_count_after / message_count_before) if message_count_before > 0 else 0.0
        self._summarization_records.append({
            "message_count_before": message_count_before,
            "message_count_after": message_count_after,
            "reduction_ratio": round(ratio, 4),
            "time_ms": time_ms,
            "timestamp": datetime.utcnow(),
        })
        self._prune_context_records(self._summarization_records)

    def record_context_filtering(self, intent: str, original_length: int, filtered_length: int) -> None:
        reduction = original_length - filtered_length
        self._filtering_records.append({
            "intent": intent,
            "original_length": original_length,
            "filtered_length": filtered_length,
            "reduction": reduction,
            "timestamp": datetime.utcnow(),
        })
        self._prune_context_records(self._filtering_records)

    def record_context_error(self, operation: str, error_type: str) -> None:
        self._context_errors.append({
            "operation": operation,
            "error_type": error_type,
            "timestamp": datetime.utcnow(),
        })
        self._prune_context_records(self._context_errors)

    def get_context_metrics_summary(self) -> Dict[str, Any]:
        rewrite_total = len(self._rewrite_records)
        method_dist: Dict[str, int] = {}
        rewrite_times = []
        cache_hits = 0
        for r in self._rewrite_records:
            method_dist[r["method"]] = method_dist.get(r["method"], 0) + 1
            rewrite_times.append(r["time_ms"])
            if r["cache_hit"]:
                cache_hits += 1

        summ_total = len(self._summarization_records)
        summ_times = [r["time_ms"] for r in self._summarization_records]
        summ_ratios = [r["reduction_ratio"] for r in self._summarization_records]

        filtering_by_intent: Dict[str, Dict[str, Any]] = {}
        for r in self._filtering_records:
            intent = r["intent"]
            if intent not in filtering_by_intent:
                filtering_by_intent[intent] = {"count": 0, "total_reduction": 0}
            filtering_by_intent[intent]["count"] += 1
            filtering_by_intent[intent]["total_reduction"] += r["reduction"]

        intent_summary = {}
        for intent, data in filtering_by_intent.items():
            avg_red = data["total_reduction"] / data["count"] if data["count"] > 0 else 0.0
            intent_summary[intent] = {"count": data["count"], "avg_reduction": round(avg_red, 2)}

        return {
            "query_rewrites": {
                "total": rewrite_total,
                "cache_hit_rate": round(cache_hits / rewrite_total, 4) if rewrite_total > 0 else 0.0,
                "avg_time_ms": round(statistics.mean(rewrite_times), 2) if rewrite_times else 0.0,
                "method_distribution": method_dist,
            },
            "summarizations": {
                "total": summ_total,
                "avg_reduction_ratio": round(statistics.mean(summ_ratios), 4) if summ_ratios else 0.0,
                "avg_time_ms": round(statistics.mean(summ_times), 2) if summ_times else 0.0,
            },
            "context_filtering": {
                "total": len(self._filtering_records),
                "by_intent": intent_summary,
            },
            "errors": {
                "total": len(self._context_errors),
                "by_operation": self._group_errors_by_operation(),
            },
        }

    def get_performance_percentiles(self, metric_name: str) -> Dict[str, float]:
        values = self._collect_metric_values(metric_name)
        if not values:
            return {"p50": 0.0, "p95": 0.0, "p99": 0.0}
        sorted_vals = sorted(values)
        n = len(sorted_vals)
        return {
            "p50": sorted_vals[int(n * 0.50)] if n > 1 else sorted_vals[0],
            "p95": sorted_vals[min(int(n * 0.95), n - 1)],
            "p99": sorted_vals[min(int(n * 0.99), n - 1)],
        }

    def export_metrics_json(self) -> str:
        data = {
            "timestamp": datetime.utcnow().isoformat(),
            "query_summary": self.get_summary(),
            "context_metrics": self.get_context_metrics_summary(),
            "uptime_seconds": round(self.uptime_seconds, 0),
        }
        return json.dumps(data, default=str)

    def get_metrics_for_dashboard(self) -> Dict[str, Any]:
        now = datetime.utcnow()
        hour_ago = now - timedelta(hours=1)
        day_ago = now - timedelta(days=1)

        def count_since(records: List[Dict], since: datetime) -> int:
            return sum(1 for r in records if r.get("timestamp", now) >= since)

        return {
            "rewrite_last_hour": count_since(self._rewrite_records, hour_ago),
            "rewrite_last_day": count_since(self._rewrite_records, day_ago),
            "summarization_last_hour": count_since(self._summarization_records, hour_ago),
            "summarization_last_day": count_since(self._summarization_records, day_ago),
            "filtering_last_hour": count_since(self._filtering_records, hour_ago),
            "filtering_last_day": count_since(self._filtering_records, day_ago),
            "errors_last_hour": count_since(self._context_errors, hour_ago),
            "errors_last_day": count_since(self._context_errors, day_ago),
            "context_summary": self.get_context_metrics_summary(),
        }

    def clear(self) -> None:
        self._records.clear()
        self._rewrite_records.clear()
        self._summarization_records.clear()
        self._filtering_records.clear()
        self._context_errors.clear()
        self._start_time = datetime.utcnow()

    def _filter_by_mode(self, mode: Optional[str]) -> List[QueryRecord]:
        if mode is None:
            return self._records
        return [r for r in self._records if r.mode == mode]

    def _group_errors_by_operation(self) -> Dict[str, int]:
        groups: Dict[str, int] = {}
        for e in self._context_errors:
            groups[e["operation"]] = groups.get(e["operation"], 0) + 1
        return groups

    def _collect_metric_values(self, metric_name: str) -> List[float]:
        if metric_name == "query_rewrite":
            return [r["time_ms"] for r in self._rewrite_records]
        if metric_name == "summarization":
            return [r["time_ms"] for r in self._summarization_records]
        if metric_name == "response_time":
            return [r.response_time for r in self._records]
        return []

    def _prune_context_records(self, records: List[Dict], retention_hours: int = 168) -> None:
        if len(records) < 50:
            return
        cutoff = datetime.utcnow() - timedelta(hours=retention_hours)
        records[:] = [r for r in records if r.get("timestamp", datetime.utcnow()) >= cutoff]


_metrics_instance: Optional[QueryMetrics] = None


def get_metrics(max_history: int = 1000) -> QueryMetrics:
    """Get or create singleton metrics instance."""
    global _metrics_instance
    
    if _metrics_instance is None:
        _metrics_instance = QueryMetrics(max_history=max_history)
    
    return _metrics_instance


def reset_metrics() -> None:
    """Reset the singleton metrics instance."""
    global _metrics_instance
    _metrics_instance = None


def record_query(
    query: str,
    mode: str,
    response_time: float,
    success: bool,
    error: Optional[str] = None,
    context_metadata: Optional[Dict[str, Any]] = None,
) -> QueryRecord:
    metrics = get_metrics()
    record = metrics.record(query, mode, response_time, success, error)
    if context_metadata:
        rewrite = context_metadata.get("rewrite")
        if rewrite and isinstance(rewrite, dict):
            metrics.record_query_rewrite(
                rewritten=rewrite.get("is_follow_up", False),
                method=rewrite.get("method", "none"),
                time_ms=rewrite.get("rewrite_time_ms", 0.0),
                cache_hit=rewrite.get("method") == "cache",
            )
        summarization = context_metadata.get("summarization")
        if summarization and isinstance(summarization, dict):
            metrics.record_conversation_summary(
                message_count_before=summarization.get("original_message_count", 0),
                message_count_after=summarization.get("summarized_message_count", 0),
                time_ms=summarization.get("time_ms", 0.0),
            )
        filtering = context_metadata.get("filtering")
        if filtering and isinstance(filtering, dict):
            metrics.record_context_filtering(
                intent=filtering.get("intent", "unknown"),
                original_length=filtering.get("original_length", 0),
                filtered_length=filtering.get("filtered_length", 0),
            )
    return record


def get_average_response_time(mode: Optional[str] = None) -> float:
    """Convenience function to get average response time."""
    metrics = get_metrics()
    return metrics.get_average_response_time(mode)


def get_query_count() -> int:
    """Convenience function to get total query count."""
    metrics = get_metrics()
    return metrics.total_queries