"""Event schema and thread-safe queue for agentradar."""
from __future__ import annotations

import queue
import threading
import uuid
from datetime import datetime, timezone
from typing import Any, Optional

# Global event queue — instrumentation writes, server reads
_event_queue: queue.Queue = queue.Queue(maxsize=10_000)
_queue_lock = threading.Lock()

# Thread-local storage for LangChain token usage captured by the LLM patch.
# Patched BaseChatModel.invoke writes here; extract_tokens reads and clears it.
_tl = threading.local()


def _capture_langchain_tokens(message: Any) -> None:
    """Store token usage from a LangChain AIMessage into thread-local."""
    try:
        usage = getattr(message, "usage_metadata", None)
        if usage and isinstance(usage, dict):
            p = usage.get("input_tokens", 0) or 0
            c = usage.get("output_tokens", 0) or 0
            t = usage.get("total_tokens", p + c) or 0
            if t > 0:
                existing = getattr(_tl, "pending_tokens", None)
                if existing:
                    # Accumulate across multiple llm.invoke calls in one agent
                    _tl.pending_tokens = {
                        "prompt": existing["prompt"] + p,
                        "completion": existing["completion"] + c,
                        "total": existing["total"] + t,
                        "model": existing.get("model"),
                    }
                else:
                    _tl.pending_tokens = {"prompt": p, "completion": c, "total": t, "model": None}
        # Also try response_metadata (some LangChain versions)
        rm = getattr(message, "response_metadata", None)
        if rm and isinstance(rm, dict):
            tu = rm.get("token_usage") or rm.get("usage")
            if isinstance(tu, dict):
                p = tu.get("prompt_tokens", 0) or tu.get("input_tokens", 0) or 0
                c = tu.get("completion_tokens", 0) or tu.get("output_tokens", 0) or 0
                t = tu.get("total_tokens", p + c) or 0
                if t > 0:
                    existing = getattr(_tl, "pending_tokens", None)
                    if not existing:
                        _tl.pending_tokens = {"prompt": p, "completion": c, "total": t, "model": None}
    except Exception:
        pass


def patch_langchain() -> None:
    """Monkey-patch BaseChatModel.invoke/ainvoke to capture token usage."""
    try:
        from langchain_core.language_models.chat_models import BaseChatModel
        if getattr(BaseChatModel, "_agentradar_patched", False):
            return

        _orig_invoke = BaseChatModel.invoke

        def _patched_invoke(self, input, config=None, **kwargs):
            result = _orig_invoke(self, input, config=config, **kwargs)
            _capture_langchain_tokens(result)
            return result

        BaseChatModel.invoke = _patched_invoke

        try:
            _orig_ainvoke = BaseChatModel.ainvoke

            async def _patched_ainvoke(self, input, config=None, **kwargs):
                result = await _orig_ainvoke(self, input, config=config, **kwargs)
                _capture_langchain_tokens(result)
                return result

            BaseChatModel.ainvoke = _patched_ainvoke
        except Exception:
            pass

        BaseChatModel._agentradar_patched = True
    except Exception:
        pass


def make_run_id() -> str:
    return f"run_{datetime.now(timezone.utc).strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:6]}"


def now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def build_event(
    event_type: str,
    run_id: str,
    agent_id: str,
    agent_name: str,
    agent_type: str = "agent",
    parent_id: Optional[str] = None,
    step_number: int = 0,
    input: Any = None,
    output: Any = None,
    state_before: Optional[dict] = None,
    state_after: Optional[dict] = None,
    tokens: Optional[dict] = None,
    duration_ms: float = 0.0,
    error: Optional[dict] = None,
    metadata: Optional[dict] = None,
) -> dict:
    return {
        "event_type": event_type,
        "run_id": run_id,
        "agent_id": agent_id,
        "agent_name": agent_name,
        "agent_type": agent_type,
        "parent_id": parent_id,
        "step_number": step_number,
        "timestamp": now_iso(),
        "input": _safe_serialize(input),
        "output": _safe_serialize(output),
        "state_before": _safe_serialize(state_before),
        "state_after": _safe_serialize(state_after),
        "tokens": tokens or {"prompt": 0, "completion": 0, "total": 0, "model": None},
        "duration_ms": duration_ms,
        "error": error,
        "metadata": metadata or {"tag": None, "notes": None},
    }


def emit(event: dict) -> None:
    """Non-blocking emit — drops event if queue is full rather than blocking pipeline."""
    try:
        _event_queue.put_nowait(event)
    except queue.Full:
        pass


def extract_tokens(output: Any, input_data: Any = None) -> dict:
    """
    Pull token usage from:
      1. Thread-local set by the LangChain patch (covers all return types)
      2. output["tokens"] dict
      3. output["tokens_<name>"] dict
      4. output with "prompt_tokens"/"completion_tokens" directly
    """
    default = {"prompt": 0, "completion": 0, "total": 0, "model": None}

    # 0. Thread-local captured by patched BaseChatModel.invoke — most reliable
    pending = getattr(_tl, "pending_tokens", None)
    if pending and (pending.get("total") or 0) > 0:
        _tl.pending_tokens = None  # consume it
        return pending

    if not isinstance(output, dict):
        return default

    # 1. Direct 'tokens' key
    t = output.get("tokens")
    if isinstance(t, dict) and (t.get("total") or 0) > 0:
        return {
            "prompt": t.get("prompt", 0),
            "completion": t.get("completion", 0),
            "total": t.get("total", 0),
            "model": t.get("model"),
        }

    # 2. tokens_<name> keys — skip ones already present in the input state
    input_keys: set = set()
    if isinstance(input_data, dict):
        # input_data is {"state": {...}, ...} — look inside the first positional arg
        first = next(iter(input_data.values()), None)
        if isinstance(first, dict):
            input_keys = set(first.keys())

    new_token_keys = [k for k in output if k.startswith("tokens_") and k not in input_keys]
    for key in reversed(new_token_keys):   # last added = most recent
        t = output[key]
        if isinstance(t, dict) and (t.get("total") or 0) > 0:
            return {
                "prompt": t.get("prompt", 0),
                "completion": t.get("completion", 0),
                "total": t.get("total", 0),
                "model": t.get("model"),
            }

    # 3. OpenAI raw style
    if "prompt_tokens" in output and "completion_tokens" in output:
        p, c = output["prompt_tokens"], output["completion_tokens"]
        return {"prompt": p, "completion": c,
                "total": output.get("total_tokens", p + c), "model": None}

    return default


def _safe_serialize(obj: Any) -> Any:
    """Convert objects to JSON-serializable form without crashing."""
    if obj is None:
        return None
    if isinstance(obj, (str, int, float, bool)):
        return obj
    if isinstance(obj, (list, tuple)):
        return [_safe_serialize(i) for i in obj]
    if isinstance(obj, dict):
        return {str(k): _safe_serialize(v) for k, v in obj.items()}
    try:
        import dataclasses
        if dataclasses.is_dataclass(obj):
            return dataclasses.asdict(obj)
    except Exception:
        pass
    try:
        return obj.model_dump()
    except Exception:
        pass
    try:
        return obj.__dict__
    except Exception:
        pass
    return str(obj)
