"""Core SDK: init, trace, span decorators, context management, batch sender."""

from __future__ import annotations

import atexit
import contextvars
import functools
import os
import secrets
import threading
import time
from datetime import datetime, timezone
from typing import Any, Callable, TypeVar

F = TypeVar("F", bound=Callable[..., Any])

# Context variables for async-safe trace propagation
_current_trace_id: contextvars.ContextVar[str | None] = contextvars.ContextVar(
    "current_trace_id", default=None
)
_current_span_id: contextvars.ContextVar[str | None] = contextvars.ContextVar(
    "current_span_id", default=None
)

# Global state
_api_key: str | None = None
_api_url: str = "https://app.veloscope.dev/api/v1/ingest"
_enabled: bool = False
_spans_buffer: list[dict[str, Any]] = []

# Batch sender state
_flush_lock = threading.Lock()
_flush_thread: threading.Thread | None = None
_shutdown_event = threading.Event()
_FLUSH_INTERVAL = 5.0   # seconds
_MAX_BUFFER_SIZE = 100  # flush immediately when buffer hits this


def _generate_trace_id() -> str:
    """Generate a 128-bit hex trace ID (32 chars)."""
    return secrets.token_hex(16)


def _generate_span_id() -> str:
    """Generate a 64-bit hex span ID (16 chars)."""
    return secrets.token_hex(8)


def _now_iso() -> str:
    """Current time in ISO 8601 format."""
    return datetime.now(timezone.utc).isoformat()


def _flush_spans() -> None:
    """Send buffered spans to the API and clear the buffer on success."""
    with _flush_lock:
        if not _spans_buffer or not _api_key:
            return
        spans_to_send = _spans_buffer.copy()
        _spans_buffer.clear()

    try:
        import httpx  # noqa: PLC0415 — intentional lazy import

        httpx.post(
            _api_url,
            json={"spans": spans_to_send},
            headers={"Authorization": f"Bearer {_api_key}"},
            timeout=5.0,
        )
    except Exception:
        # Graceful degradation: silently drop on any network/server error
        pass


def flush() -> None:
    """Manually flush all buffered spans to the API."""
    _flush_spans()


def _flush_worker() -> None:
    """Background daemon thread: flush on interval."""
    while not _shutdown_event.wait(_FLUSH_INTERVAL):
        if _enabled:
            _flush_spans()


def _start_flush_thread() -> None:
    global _flush_thread
    if _flush_thread is None or not _flush_thread.is_alive():
        _shutdown_event.clear()
        _flush_thread = threading.Thread(
            target=_flush_worker,
            daemon=True,
            name="veloscope-flush",
        )
        _flush_thread.start()
        atexit.register(_flush_spans)


def init(
    api_key: str | None = None,
    api_url: str | None = None,
    auto_patch: bool = True,
) -> None:
    """Initialize the Veloscope SDK.

    Args:
        api_key: Your Veloscope API key. Falls back to VELOSCOPE_API_KEY env var.
        api_url: Custom API endpoint URL. Defaults to Veloscope cloud.
        auto_patch: Whether to auto-patch openai/anthropic. Defaults to True.
    """
    global _api_key, _api_url, _enabled

    _api_key = api_key or os.environ.get("VELOSCOPE_API_KEY")
    if not _api_key:
        # Graceful degradation: no key = no-op
        return

    if api_url:
        _api_url = api_url

    _enabled = True
    _start_flush_thread()

    if auto_patch:
        _auto_patch_providers()


def _auto_patch_providers() -> None:
    """Auto-patch supported LLM providers."""
    # TODO: Implement openai auto-patcher
    # TODO: Implement anthropic auto-patcher
    pass


def trace(name: str | None = None) -> Callable[[F], F]:
    """Decorator to create a trace (top-level agent run).

    Usage:
        @trace("my-agent")
        def run_agent(query: str):
            ...
    """

    def decorator(func: F) -> F:
        trace_name = name or func.__name__

        @functools.wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            if not _enabled:
                return func(*args, **kwargs)

            trace_id = _generate_trace_id()
            span_id = _generate_span_id()

            token_trace = _current_trace_id.set(trace_id)
            token_span = _current_span_id.set(span_id)

            start = time.perf_counter()
            start_time = _now_iso()
            status = "ok"
            error_msg = None

            try:
                result = func(*args, **kwargs)
                return result
            except Exception as e:
                status = "error"
                error_msg = str(e)
                raise
            finally:
                end_time = _now_iso()
                duration_ms = int((time.perf_counter() - start) * 1000)

                _spans_buffer.append(
                    {
                        "trace_id": trace_id,
                        "span_id": span_id,
                        "parent_span_id": None,
                        "name": trace_name,
                        "span_type": "custom",
                        "status": status,
                        "start_time": start_time,
                        "end_time": end_time,
                        "duration_ms": duration_ms,
                        "error_message": error_msg,
                    }
                )

                _current_trace_id.reset(token_trace)
                _current_span_id.reset(token_span)

                if len(_spans_buffer) >= _MAX_BUFFER_SIZE:
                    threading.Thread(
                        target=_flush_spans, daemon=True, name="veloscope-flush-eager"
                    ).start()

        return wrapper  # type: ignore[return-value]

    return decorator


def span(name: str | None = None, span_type: str = "custom") -> Callable[[F], F]:
    """Decorator to create a span within a trace.

    Usage:
        @span("retrieve-docs", span_type="tool")
        def retrieve(query: str):
            ...
    """

    def decorator(func: F) -> F:
        span_name = name or func.__name__

        @functools.wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            if not _enabled:
                return func(*args, **kwargs)

            trace_id = _current_trace_id.get()
            if not trace_id:
                # No active trace, run without instrumentation
                return func(*args, **kwargs)

            parent_span_id = _current_span_id.get()
            new_span_id = _generate_span_id()
            token_span = _current_span_id.set(new_span_id)

            start = time.perf_counter()
            start_time = _now_iso()
            status = "ok"
            error_msg = None

            try:
                result = func(*args, **kwargs)
                return result
            except Exception as e:
                status = "error"
                error_msg = str(e)
                raise
            finally:
                end_time = _now_iso()
                duration_ms = int((time.perf_counter() - start) * 1000)

                _spans_buffer.append(
                    {
                        "trace_id": trace_id,
                        "span_id": new_span_id,
                        "parent_span_id": parent_span_id,
                        "name": span_name,
                        "span_type": span_type,
                        "status": status,
                        "start_time": start_time,
                        "end_time": end_time,
                        "duration_ms": duration_ms,
                        "error_message": error_msg,
                    }
                )

                _current_span_id.reset(token_span)

        return wrapper  # type: ignore[return-value]

    return decorator
