from __future__ import annotations

from collections import deque
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from threading import Lock
from typing import Deque, Dict, List, Optional, Tuple

from .models import TaskMetrics


@dataclass
class _TaskRecord:
    task_id: str
    task_type: str
    started_at: datetime
    ended_at: Optional[datetime] = None
    execution_time_ms: Optional[int] = None
    success: Optional[bool] = None
    error_type: Optional[str] = None
    error_message: Optional[str] = None


class MetricsCollector:
    """
    In-memory rolling metrics with simple window filtering.
    Not persistent; intended for live agent reporting.
    """

    def __init__(self, max_records: int = 5000):
        self.records: Deque[_TaskRecord] = deque(maxlen=max_records)
        self._lock = Lock()

    def record_start(self, task_id: str, task_type: str, ts: Optional[datetime] = None) -> None:
        with self._lock:
            self.records.append(
                _TaskRecord(task_id=task_id, task_type=task_type, started_at=ts or datetime.now(timezone.utc))
            )

    def _find_record(self, task_id: str) -> Optional[_TaskRecord]:
        for rec in reversed(self.records):
            if rec.task_id == task_id:
                return rec
        return None

    def record_success(self, task_id: str, execution_time_ms: int, ts: Optional[datetime] = None) -> None:
        with self._lock:
            rec = self._find_record(task_id)
            if not rec:
                rec = _TaskRecord(task_id=task_id, task_type="unknown", started_at=ts or datetime.now(timezone.utc))
                self.records.append(rec)
            rec.success = True
            rec.execution_time_ms = execution_time_ms
            rec.ended_at = ts or datetime.now(timezone.utc)

    def record_failure(self, task_id: str, error_type: str, error_message: str, ts: Optional[datetime] = None) -> None:
        with self._lock:
            rec = self._find_record(task_id)
            if not rec:
                rec = _TaskRecord(task_id=task_id, task_type="unknown", started_at=ts or datetime.now(timezone.utc))
                self.records.append(rec)
            rec.success = False
            rec.error_type = error_type
            rec.error_message = error_message
            rec.ended_at = ts or datetime.now(timezone.utc)

    def _filter_window(self, window: timedelta) -> List[_TaskRecord]:
        cutoff = datetime.now(timezone.utc) - window
        return [r for r in self.records if r.started_at >= cutoff]

    def _latency_percentile(self, latencies: List[int], pct: float) -> float:
        if not latencies:
            return 0.0
        k = max(min(int(len(latencies) * pct) - 1, len(latencies) - 1), 0)
        return float(sorted(latencies)[k])

    def snapshot(self, window: timedelta) -> TaskMetrics:
        with self._lock:
            window_records = self._filter_window(window)
            total = len(window_records)
            successes = sum(1 for r in window_records if r.success)
            failures = sum(1 for r in window_records if r.success is False)
            timeout_errors = sum(1 for r in window_records if r.error_type == "timeout")
            llm_errors = sum(1 for r in window_records if r.error_type == "llm_error")
            db_errors = sum(1 for r in window_records if r.error_type == "db_error")
            desensitization_errors = sum(1 for r in window_records if r.error_type == "desensitization_error")
            latencies = [r.execution_time_ms for r in window_records if r.execution_time_ms is not None]

            avg_latency = float(sum(latencies) / len(latencies)) if latencies else 0.0
            p95 = self._latency_percentile(latencies, 0.95) if latencies else 0.0
            p99 = self._latency_percentile(latencies, 0.99) if latencies else 0.0

            return TaskMetrics(
                total_tasks=total,
                successful_tasks=successes,
                failed_tasks=failures,
                timeout_errors=timeout_errors,
                llm_errors=llm_errors,
                db_errors=db_errors,
                desensitization_errors=desensitization_errors,
                avg_execution_time_ms=avg_latency,
                p95_execution_time_ms=p95,
                p99_execution_time_ms=p99,
            )
