from __future__ import annotations

from contextvars import ContextVar
from datetime import datetime, timezone
from typing import Any
from uuid import UUID, uuid4

from agenttrace.client import IngestClient
from agenttrace.contracts import Run, Span, SpanType, build_ingest_payload

_current_run_context: ContextVar[RunTraceContext | None] = ContextVar(
    "agenttrace_current_run_context", default=None
)


def _now_utc() -> datetime:
    return datetime.now(timezone.utc)


class RunTraceContext:
    def __init__(self, name: str, metadata: dict[str, Any] | None = None) -> None:
        self._run = Run(run_id=uuid4(), name=name, created_at=_now_utc(), metadata_json=metadata or {})
        self._spans: list[Span] = []
        self._entered = False
        self._token = None
        self._flushed = False
        self._flush_error: Exception | None = None

    @property
    def run_id(self) -> str:
        return str(self._run.run_id)

    @property
    def spans(self) -> list[Span]:
        return list(self._spans)

    @property
    def flush_error(self) -> Exception | None:
        return self._flush_error

    def add_span(
        self,
        *,
        span_type: SpanType,
        name: str,
        start_time: datetime,
        end_time: datetime,
        input_json: Any,
        output_json: Any,
        error: str | None = None,
        metadata_json: dict[str, Any] | None = None,
        parent_span_id: str | None = None,
        span_id: str | None = None,
    ) -> Span:
        span_uuid = UUID(span_id) if span_id else uuid4()
        parent_uuid = None
        if parent_span_id:
            matching_parent = next((s.span_id for s in self._spans if str(s.span_id) == parent_span_id), None)
            parent_uuid = matching_parent or UUID(parent_span_id)

        duration_ms = int((end_time - start_time).total_seconds() * 1000)
        span = Span(
            span_id=span_uuid,
            run_id=self._run.run_id,
            parent_span_id=parent_uuid,
            span_type=span_type,
            name=name,
            start_time=start_time,
            end_time=end_time,
            duration_ms=duration_ms,
            input_json=input_json,
            output_json=output_json,
            error=error,
            metadata_json=metadata_json or {},
        )
        span.validate()
        self._spans.append(span)
        return span

    def build_payload(self) -> dict[str, Any]:
        return build_ingest_payload([self._run], self._spans)

    def flush(self) -> None:
        if self._flushed:
            return
        client = IngestClient.from_env()
        if client is None:
            self._flushed = True
            return
        payload = self.build_payload()
        client.ingest(payload)
        self._flushed = True

    def __enter__(self) -> RunTraceContext:
        if self._entered:
            raise RuntimeError("RunTraceContext cannot be entered twice")
        self._token = _current_run_context.set(self)
        self._entered = True
        return self

    def __exit__(self, exc_type, exc, tb) -> bool:
        if exc is not None:
            now = _now_utc()
            self.add_span(
                span_type="agent",
                name="run_error",
                start_time=now,
                end_time=now,
                input_json={},
                output_json={},
                error=str(exc),
                metadata_json={},
            )
        try:
            self.flush()
        except Exception as flush_error:  # noqa: BLE001
            self._flush_error = flush_error
        finally:
            if self._token is not None:
                _current_run_context.reset(self._token)
        return False

    async def aflush(self) -> None:
        if self._flushed:
            return
        client = IngestClient.from_env()
        if client is None:
            self._flushed = True
            return
        payload = self.build_payload()
        await client.async_ingest(payload)
        self._flushed = True

    async def __aenter__(self) -> RunTraceContext:
        if self._entered:
            raise RuntimeError("RunTraceContext cannot be entered twice")
        self._token = _current_run_context.set(self)
        self._entered = True
        return self

    async def __aexit__(self, exc_type, exc, tb) -> bool:
        if exc is not None:
            now = _now_utc()
            self.add_span(
                span_type="agent",
                name="run_error",
                start_time=now,
                end_time=now,
                input_json={},
                output_json={},
                error=str(exc),
                metadata_json={},
            )
        try:
            await self.aflush()
        except Exception as flush_error:  # noqa: BLE001
            self._flush_error = flush_error
        finally:
            if self._token is not None:
                _current_run_context.reset(self._token)
        return False


class TraceAPI:
    def run(self, name: str, metadata: dict[str, Any] | None = None) -> RunTraceContext:
        return RunTraceContext(name=name, metadata=metadata)

    def current(self) -> RunTraceContext | None:
        return _current_run_context.get()


trace = TraceAPI()
