"""Cost Telemetry Collector - Module 1."""

import json
import os
from typing import List, Dict, Any, Optional
from datetime import datetime
from pathlib import Path

from .trace_schema import AgentTrace, TraceStep, ModelCall, ToolCall, VerifierCall, TaskType, Outcome, FailureTag


class CostTelemetryCollector:
    """Collects structured telemetry from agent runs and persists normalized traces."""

    def __init__(self, storage_path: str = "./traces"):
        self.storage_path = Path(storage_path)
        self.storage_path.mkdir(parents=True, exist_ok=True)
        self._pending: Dict[str, AgentTrace] = {}

    def start_trace(self, trace_id: str, user_request: str, task_type: TaskType) -> AgentTrace:
        trace = AgentTrace(
            trace_id=trace_id,
            user_request=user_request,
            task_type=task_type,
        )
        self._pending[trace_id] = trace
        return trace

    def add_step(
        self,
        trace_id: str,
        step_id: str,
        model_call: ModelCall,
        tool_calls: Optional[List[ToolCall]] = None,
        verifier_calls: Optional[List[VerifierCall]] = None,
        context_size_tokens: int = 0,
        context_sources: Optional[List[str]] = None,
        retry_count: int = 0,
        recovery_action: Optional[str] = None,
        artifacts_created: Optional[List[str]] = None,
        step_outcome: Optional[Outcome] = None,
    ) -> None:
        trace = self._pending.get(trace_id)
        if not trace:
            raise ValueError(f"Trace {trace_id} not found")
        
        step = TraceStep(
            step_id=step_id,
            timestamp=datetime.utcnow(),
            task_type=trace.task_type,
            model_call=model_call,
            tool_calls=tool_calls or [],
            verifier_calls=verifier_calls or [],
            context_size_tokens=context_size_tokens,
            context_sources=context_sources or [],
            retry_count=retry_count,
            recovery_action=recovery_action,
            artifacts_created=artifacts_created or [],
            step_outcome=step_outcome,
        )
        trace.steps.append(step)

    def finalize_trace(
        self,
        trace_id: str,
        final_outcome: Outcome,
        failure_tags: Optional[List[FailureTag]] = None,
        user_satisfaction: Optional[float] = None,
        total_cost_saved_vs_frontier: Optional[float] = None,
        optimal_cost: Optional[float] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> AgentTrace:
        trace = self._pending.pop(trace_id)
        trace.final_outcome = final_outcome
        trace.failure_tags = failure_tags or []
        trace.user_satisfaction = user_satisfaction
        trace.total_cost_saved_vs_frontier = total_cost_saved_vs_frontier
        trace.optimal_cost = optimal_cost
        trace.metadata = metadata or {}
        trace.total_cost = trace.total_cost_computed
        
        self._persist(trace)
        return trace

    def _persist(self, trace: AgentTrace) -> None:
        filepath = self.storage_path / f"{trace.trace_id}.json"
        with open(filepath, "w") as f:
            json.dump(trace.to_dict(), f, indent=2, default=str)

    def load_trace(self, trace_id: str) -> Optional[AgentTrace]:
        filepath = self.storage_path / f"{trace_id}.json"
        if not filepath.exists():
            return None
        with open(filepath, "r") as f:
            data = json.load(f)
        # Simplified deserialization - full version would reconstruct dataclasses
        return data

    def list_traces(self) -> List[str]:
        return [p.stem for p in self.storage_path.glob("*.json")]

    def get_stats(self) -> Dict[str, Any]:
        traces = []
        for tid in self.list_traces():
            t = self.load_trace(tid)
            if t:
                traces.append(t)
        
        if not traces:
            return {"count": 0}
        
        total_cost = sum(t.get("total_cost", 0) for t in traces if isinstance(t, dict))
        total_steps = sum(len(t.get("steps", [])) for t in traces if isinstance(t, dict))
        
        return {
            "count": len(traces),
            "avg_cost": total_cost / len(traces),
            "avg_steps": total_steps / len(traces),
            "success_rate": sum(
                1 for t in traces
                if isinstance(t, dict) and t.get("final_outcome") == "success"
            ) / len(traces),
        }