import time
import uuid
import statistics
from typing import Any

TRACE_TTL = 3600


class TraceEvent:
    def __init__(self, event_type: str, data: dict):
        self.id: str = str(uuid.uuid4())
        self.type: str = event_type
        self.data: dict = data
        self.start_time: float = time.time()
        self.end_time: float | None = None
        self.steps: list[dict] = []
        self.success: bool = True


class ObservabilityBrain:
    def __init__(self):
        self._traces: dict[str, TraceEvent] = {}
        self._metrics: dict[str, list[float]] = {
            "request_latency": [],
            "policy_latency": [],
            "memory_latency": [],
            "service_latency": [],
        }
        self._error_count: int = 0
        self._request_count: int = 0

    def _cleanup_old_traces(self) -> None:
        cutoff = time.time() - TRACE_TTL
        expired = [tid for tid, t in self._traces.items() if (t.end_time or time.time()) < cutoff]
        for tid in expired:
            del self._traces[tid]

    def start_trace(self, event: dict) -> str:
        if len(self._traces) > 10000:
            self._cleanup_old_traces()
        trace = TraceEvent(event.get("type", "unknown"), event)
        self._traces[trace.id] = trace
        self._request_count += 1
        return trace.id

    def add_step(self, trace_id: str, step_name: str, data: Any = None) -> None:
        trace = self._traces.get(trace_id)
        if trace is None:
            return
        trace.steps.append({
            "step": step_name,
            "data": data,
            "timestamp": time.time(),
        })

    def record_latency(self, metric_name: str, value: float) -> None:
        if metric_name in self._metrics:
            self._metrics[metric_name].append(value)
            if len(self._metrics[metric_name]) > 1000:
                self._metrics[metric_name] = self._metrics[metric_name][-1000:]

    def end_trace(self, trace_id: str, success: bool = True) -> None:
        trace = self._traces.get(trace_id)
        if trace is None:
            return
        trace.end_time = time.time()
        trace.success = success
        if not success:
            self._error_count += 1
        latency = trace.end_time - trace.start_time
        self.record_latency("request_latency", latency)

    def get_trace(self, trace_id: str) -> dict | None:
        trace = self._traces.get(trace_id)
        if trace is None:
            return None
        return {
            "id": trace.id,
            "type": trace.type,
            "duration": (trace.end_time or time.time()) - trace.start_time,
            "success": trace.success,
            "steps": trace.steps,
        }

    def replay(self, trace_id: str) -> dict | None:
        return self.get_trace(trace_id)

    def get_metrics(self) -> dict:
        def _avg(values: list[float]) -> float:
            return round(statistics.mean(values), 3) if values else 0.0

        def _p95(values: list[float]) -> float:
            if not values:
                return 0.0
            sorted_vals = sorted(values)
            idx = int(len(sorted_vals) * 0.95)
            return round(sorted_vals[min(idx, len(sorted_vals) - 1)], 3)

        return {
            "requests": self._request_count,
            "errors": self._error_count,
            "error_rate": round(self._error_count / max(self._request_count, 1), 4),
            "avg_latency": _avg(self._metrics["request_latency"]),
            "p95_latency": _p95(self._metrics["request_latency"]),
            "avg_policy_latency": _avg(self._metrics["policy_latency"]),
            "avg_memory_latency": _avg(self._metrics["memory_latency"]),
            "avg_service_latency": _avg(self._metrics["service_latency"]),
            "active_traces": sum(
                1 for t in self._traces.values() if t.end_time is None
            ),
        }

    def reset(self) -> None:
        self._traces.clear()
        for k in self._metrics:
            self._metrics[k].clear()
        self._error_count = 0
        self._request_count = 0