import contextvars
import time
import uuid
from contextlib import contextmanager
from typing import Optional

from .client import EventClient
from .event import CostEvent
from .pricing import PricingDB

# Contextvar so @track_tool decorators can find the active run
_active_run: contextvars.ContextVar[Optional["RunTracker"]] = contextvars.ContextVar(
    "_active_run", default=None
)


class RunTracker:
    """Tracks costs within a single agent run (task execution)."""

    def __init__(
        self,
        client: EventClient,
        pricing: PricingDB,
        agent_id: str,
        run_id: str,
        metadata: dict,
    ):
        self._client = client
        self._pricing = pricing
        self._agent_id = agent_id
        self._run_id = run_id
        self._metadata = metadata
        self._events: list[CostEvent] = []

    def record(
        self,
        component_name: str,
        *,
        cost: Optional[float] = None,
        event_type: str = "tool_call",
        duration_ms: Optional[int] = None,
        tokens_in: Optional[int] = None,
        tokens_out: Optional[int] = None,
        metadata: Optional[dict] = None,
    ) -> CostEvent:
        """Record a cost event manually."""
        # Resolve cost from pricing DB if not provided
        if cost is None:
            if event_type == "llm_call" and tokens_in is not None:
                cost = self._pricing.llm_cost(
                    component_name, tokens_in, tokens_out or 0
                )
            else:
                cost = self._pricing.tool_cost(component_name)

        merged_metadata = {**self._metadata, **(metadata or {})}

        event = CostEvent(
            agent_id=self._agent_id,
            run_id=self._run_id,
            component_name=component_name,
            event_type=event_type,
            cost_usd=cost,
            duration_ms=duration_ms,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            metadata=merged_metadata,
        )

        self._events.append(event)
        self._client.enqueue(event)
        return event

    def record_llm(
        self,
        model: str,
        tokens_in: int,
        tokens_out: int,
        *,
        duration_ms: Optional[int] = None,
        cost: Optional[float] = None,
        metadata: Optional[dict] = None,
    ) -> CostEvent:
        """Convenience method for recording LLM calls."""
        return self.record(
            model,
            event_type="llm_call",
            cost=cost,
            duration_ms=duration_ms,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            metadata=metadata,
        )

    @property
    def total_cost(self) -> float:
        return sum(e.cost_usd for e in self._events)

    @property
    def events(self) -> list[CostEvent]:
        return list(self._events)


class Tracker:
    """Main tracker that creates run contexts for an agent."""

    def __init__(
        self,
        client: EventClient,
        pricing: PricingDB,
        agent_id: str,
        default_metadata: Optional[dict] = None,
    ):
        self._client = client
        self._pricing = pricing
        self._agent_id = agent_id
        self._default_metadata = default_metadata or {}

    @contextmanager
    def run(self, task: Optional[str] = None, **extra_metadata):
        """Context manager that tracks a single agent run."""
        run_id = str(uuid.uuid4())
        metadata = {**self._default_metadata, **extra_metadata}
        if task:
            metadata["task"] = task

        run_tracker = RunTracker(
            client=self._client,
            pricing=self._pricing,
            agent_id=self._agent_id,
            run_id=run_id,
            metadata=metadata,
        )

        # Set this run as the active context for @track_tool decorators
        token = _active_run.set(run_tracker)
        start_time = time.time()
        try:
            yield run_tracker
        finally:
            elapsed_ms = int((time.time() - start_time) * 1000)
            _active_run.reset(token)

            # Send a run summary event
            summary = CostEvent(
                agent_id=self._agent_id,
                run_id=run_id,
                component_name="run_summary",
                event_type="run_end",
                cost_usd=run_tracker.total_cost,
                duration_ms=elapsed_ms,
                metadata={
                    **metadata,
                    "event_count": len(run_tracker.events),
                },
            )
            self._client.enqueue(summary)
