from __future__ import annotations

import atexit
import asyncio
import contextvars
import threading
from contextvars import ContextVar
from dataclasses import fields as dataclass_fields
from datetime import datetime, timezone
from functools import wraps
from types import TracebackType
from typing import Any, Callable, Coroutine, TypeVar

from .models import Span, Trace


# --- Sync flush registry ---------------------------------------------------
# Background sync threads register here so atexit can wait for them to finish.
# Without this, daemon threads are killed on process exit and syncs are lost.

_pending_syncs: list[threading.Thread] = []
_sync_lock = threading.Lock()
_FLUSH_TIMEOUT = 5.0  # seconds to wait per thread at exit


def _register_sync(thread: threading.Thread) -> None:
    with _sync_lock:
        # Prune completed threads to avoid unbounded growth
        _pending_syncs[:] = [t for t in _pending_syncs if t.is_alive()]
        _pending_syncs.append(thread)


def _flush_syncs() -> None:
    """Wait for pending sync threads to complete (called at process exit)."""
    with _sync_lock:
        threads = list(_pending_syncs)
        _pending_syncs.clear()
    for t in threads:
        t.join(timeout=_FLUSH_TIMEOUT)


def flush() -> None:
    """Block until all pending background syncs complete.

    Call this explicitly in short-lived processes (Lambda, scripts) to
    ensure trace summaries are delivered before the process exits.
    For long-running processes, atexit handles this automatically.
    """
    _flush_syncs()


atexit.register(_flush_syncs)

_current_trace: ContextVar[TraceContext | None] = ContextVar("_current_trace", default=None)
_current_span: ContextVar[SpanContext | None] = ContextVar("_current_span", default=None)

T = TypeVar("T")


def get_current_trace() -> TraceContext | None:
    """Get the current trace context, if any."""
    return _current_trace.get()


def get_current_span() -> SpanContext | None:
    """Get the current span context, if any."""
    return _current_span.get()


def spawn_task(coro: Coroutine[Any, Any, T]) -> asyncio.Task[T]:
    """
    Spawn an async task that inherits the current span context.

    This is a drop-in replacement for asyncio.create_task() that ensures
    child tasks see the parent's span hierarchy.

    Example:
        async def worker():
            with t.span("work", "tool") as s:
                await s.do_work()

        # Use spawn_task instead of asyncio.create_task()
        task = spawn_task(worker())
        await task

    Without this, child tasks created via asyncio.create_task() won't
    see the parent's spans because ContextVars are task-local.
    """
    ctx = contextvars.copy_context()
    return asyncio.create_task(coro, context=ctx)


def run_in_context(coro: Coroutine[Any, Any, T]) -> T:
    """
    Run an awaitable in the current span context.

    Use this when you need to run async code synchronously (e.g., in a
    thread pool) while preserving span context.

    Example:
        async def async_work():
            with t.span("work", "tool"):
                return await do_stuff()

        result = run_in_context(async_work())
    """
    ctx = contextvars.copy_context()

    async def _run():
        return await asyncio.create_task(coro, context=ctx)

    return asyncio.run(_run())

# Fields that can be set via done() kwargs — excludes identifiers and timestamps
_SPAN_SETTABLE = frozenset(
    f.name for f in dataclass_fields(Span)
    if f.name not in ("span_id", "parent_id", "name", "kind", "started_at")
)
_TRACE_SETTABLE = frozenset(
    f.name for f in dataclass_fields(Trace)
    if f.name not in ("trace_id", "project", "started_at", "spans")
)


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


class SpanContext:
    def __init__(
        self,
        trace_ctx: TraceContext,
        name: str,
        kind: str,
        **kwargs: Any,
    ) -> None:
        parent = _current_span.get()
        parent_id = parent.span.span_id if parent is not None else None
        self.span = Span(name=name, kind=kind, parent_id=parent_id, **kwargs)
        self._trace_ctx = trace_ctx
        self._token: Any = None

    def done(self, *, output: Any = None, **kwargs: Any) -> None:
        if output is not None:
            self.span.output = output
        for k, v in kwargs.items():
            if k not in _SPAN_SETTABLE:
                raise ValueError(f"Cannot set '{k}' on Span via done()")
            setattr(self.span, k, v)

    def __enter__(self) -> SpanContext:
        self._token = _current_span.set(self)
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        self.span.ended_at = _now()
        if exc_val is not None:
            self.span.status = "error"
            self.span.error = f"{type(exc_val).__name__}: {exc_val}"
        self._trace_ctx.trace.spans.append(self.span)
        _current_span.reset(self._token)


class TraceContext:
    def __init__(
        self,
        project: str,
        *,
        writer: Any = None,
        **kwargs: Any,
    ) -> None:
        self.trace = Trace(project=project, **kwargs)
        self._writer = writer
        self._token: Any = None

    @property
    def trace_id(self) -> str:
        return self.trace.trace_id

    def span(self, name: str, kind: str, **kwargs: Any) -> SpanContext:
        return SpanContext(self, name, kind, **kwargs)

    def done(self, result: Any = None, *, output: Any = None, **kwargs: Any) -> None:
        # If result is a framework object, try auto-extraction
        if result is not None:
            from .mappers import extract

            extracted = extract(result)
            if extracted is not None:
                mapped_output, spans = extracted
                if output is None:
                    self.trace.output = mapped_output
                # Skip mapper LLM spans if the trace already has LLM spans
                # with model data (user created them via explicit span contexts)
                existing_models = {
                    s.model for s in self.trace.spans if s.model is not None
                }
                if existing_models:
                    spans = [s for s in spans if s.kind != "llm"]
                self.trace.spans.extend(spans)
            elif output is None:
                # Not a recognized framework result — treat as raw output
                self.trace.output = result

        if output is not None:
            self.trace.output = output
        for k, v in kwargs.items():
            if k not in _TRACE_SETTABLE:
                raise ValueError(f"Cannot set '{k}' on Trace via done()")
            setattr(self.trace, k, v)

    def __enter__(self) -> TraceContext:
        self._token = _current_trace.set(self)
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        self.trace.ended_at = _now()
        if exc_val is not None:
            self.trace.status = "error"
            self.trace.error = f"{type(exc_val).__name__}: {exc_val}"
        
        # Write to storage if writer is configured
        if self._writer is not None:
            self._writer.write(self.trace)
        
        # Fire-and-forget sync to API (runs regardless of storage)
        self._sync_to_api()
        
        _current_trace.reset(self._token)

    def _sync_to_api(self) -> None:
        """Sync layer 2 summary to Bandito API (fire-and-forget by default).

        Set BANDITO_SYNC_MODE=sync for synchronous behavior (blocks until complete).
        """
        import logging
        import os
        import threading

        logger = logging.getLogger(__name__)

        def _do_sync():
            try:
                from .config import load_config, resolve_sync_enabled
                from .summary import compute_layer2_from_trace_model
                from .cli.api import BanditoClient

                config = load_config()
                if not resolve_sync_enabled(config):
                    return

                client = BanditoClient()
                summary = compute_layer2_from_trace_model(self.trace)
                client.sync_summaries([summary])
            except Exception as e:
                logger.warning(f"Failed to sync summary to API: {e}")

        sync_mode = os.environ.get("BANDITO_SYNC_MODE", "async")
        if sync_mode == "sync":
            _do_sync()
        else:
            t = threading.Thread(target=_do_sync, daemon=True)
            t.start()
            _register_sync(t)
