"""@radar decorator — explicit instrumentation for custom pipelines."""
from __future__ import annotations

import asyncio
import functools
import inspect
import time
import traceback
import uuid
from typing import Any, Callable, Optional

from agentradar.server.events import build_event, emit, make_run_id, now_iso, extract_tokens
import threading

# Module-level run tracking
_current_run_id: Optional[str] = None
_step_counter: int = 0
_state_lock = threading.Lock()

# Kept for profiler compatibility (profiler reads/writes these)
_last_completed_id: Optional[str] = None
_active_count: int = 0

# Registry of original function code objects wrapped by @radar.
# The profiler checks this to avoid double-tracing these functions.
_decorated_code_objects: set = set()

# Tag registry: tag → list of agent_ids that ran with that tag.
# Populated at agent START so llm_input list resolution works even
# while the parent agent is still running (e.g. during auto-gather).
_tag_to_agent_ids: dict = {}  # tag (str) → [agent_id, ...]

# Tag → callable registry: tag → wrapped callable.
_tag_to_callable: dict = {}  # tag (str) → wrapped callable

# Tag → llm_input at decoration time: used to identify coordinators (llm_input=None).
# Only tags with llm_input=None are coordinators eligible for auto-spawn.
_tag_llm_input: dict = {}  # tag (str) → llm_input value at decoration time

# Reverse map: parent_tag → [child_callables that declared llm_input=[parent_tag, ...]].
# Built at decoration time. When a coordinator with tag="X" runs, the library looks up
# _children_of_tag["X"] and auto-gathers all of them in parallel.
_children_of_tag: dict = {}  # parent_tag (str) → [wrapped callable, ...]


def _ensure_run() -> str:
    global _current_run_id, _step_counter
    if _current_run_id is None:
        _current_run_id = make_run_id()
        _step_counter = 0
        emit(build_event(
            event_type="run_start",
            run_id=_current_run_id,
            agent_id="__run__",
            agent_name="run",
            agent_type="agent",
        ))
    return _current_run_id


def _next_step() -> int:
    global _step_counter
    _step_counter += 1
    return _step_counter


def reset_run() -> None:
    """Call this to start a fresh run."""
    global _current_run_id, _step_counter, _last_completed_id, _active_count
    if _current_run_id:
        emit(build_event(
            event_type="run_end",
            run_id=_current_run_id,
            agent_id="__run__",
            agent_name="run",
            agent_type="agent",
        ))
    _current_run_id = None
    _step_counter = 0
    _last_completed_id = None
    _active_count = 0
    _tag_to_agent_ids.clear()
    # _tag_to_callable and _children_of_tag are NOT cleared — they're built
    # at decoration time and persist across runs.


def get_current_run_id() -> Optional[str]:
    return _current_run_id


def _resolve_llm_input_parents(llm_input: Any) -> Optional[Any]:
    """
    If llm_input is a list of tag names, resolve each to ALL agent_ids that
    ran with that tag (a tag may run multiple times, e.g. proponent round 1 + 2).
    Returns a flat list of all resolved agent_ids, or None if nothing resolved.
    """
    if not isinstance(llm_input, list) or not llm_input:
        return None
    resolved = []
    with _state_lock:
        for tag_name in llm_input:
            ids = _tag_to_agent_ids.get(tag_name)
            if ids:
                resolved.extend(ids)  # ALL runs of this tag, not just the last
    if not resolved:
        return None
    return resolved if len(resolved) > 1 else resolved[0]


class radar:
    """
    Decorator for tracing agent functions.

    EDGE RULE:
      - An agent with llm_input=["parent_tag"] declares its upstream.
        The library draws an edge parent→this and auto-gathers the agent
        as a child of the parent coordinator when the parent runs.
      - An agent with no llm_input (or a string) is a root node — no
        incoming edges, no auto-spawn.

    COORDINATOR PATTERN (fan-out):
      Define coordinator with no llm_input (it's the root).
      Define children with llm_input=["coordinator_tag"].
      When coordinator runs, the library sees all registered children and
      automatically spawns them in parallel — no asyncio.gather needed.

        @radar(tag="coordinator")
        async def coordinator(query): ...   # library auto-spawns children

        @radar(tag="fetch_a", llm_input=["coordinator"])
        async def fetch_a(query): ...       # declares coordinator as parent

        @radar(tag="fetch_b", llm_input=["coordinator"])
        async def fetch_b(query): ...

    MERGE PATTERN (fan-in):
      Aggregator declares multiple parents → edges from all of them:

        @radar(tag="aggregator", llm_input=["fetch_a", "fetch_b", "fetch_c"])
        async def aggregator(results): ...

    SEQUENTIAL PATTERN:
      Each agent declares its single upstream:

        @radar(tag="writer", llm_input=["researcher"])
        def writer(notes): ...
    """

    def __new__(cls, func=None, *, tag: Optional[str] = None, notes: Optional[str] = None, llm_input: Optional[Any] = None):
        if func is not None:
            # @radar without parentheses
            instance = object.__new__(cls)
            instance._tag = None
            instance._notes = None
            instance._llm_input = None
            return instance._wrap(func)
        # @radar(...) with parentheses
        instance = object.__new__(cls)
        instance._tag = tag
        instance._notes = notes
        instance._llm_input = llm_input
        return instance

    def __call__(self, func: Callable) -> Callable:
        return self._wrap(func)

    def _wrap(self, func: Callable) -> Callable:
        tag = getattr(self, "_tag", None)
        notes = getattr(self, "_notes", None)
        llm_input = getattr(self, "_llm_input", None)
        agent_name = func.__name__
        _decorated_code_objects.add(func.__code__)

        if asyncio.iscoroutinefunction(func):
            @functools.wraps(func)
            async def async_wrapper(*args, **kwargs):
                agent_id = f"{agent_name}_{uuid.uuid4().hex[:8]}"
                return await _execute_async(
                    func, args, kwargs, agent_name, agent_id, tag, notes, llm_input
                )
            async_wrapper.__wrapped_tag__ = tag
            if tag:
                _tag_to_callable[tag] = async_wrapper
                _tag_llm_input[tag] = llm_input
            # Auto-spawn: register as child ONLY when parent is a coordinator (llm_input=None)
            # and this node has exactly one parent. Sequential nodes (with their own llm_input)
            # are called manually — auto-spawning them would pass the wrong args.
            if isinstance(llm_input, list) and len(llm_input) == 1:
                parent_tag = llm_input[0]
                parent_llm_input = _tag_llm_input.get(parent_tag, "UNKNOWN")
                if parent_llm_input is None:
                    _children_of_tag.setdefault(parent_tag, [])
                    if async_wrapper not in _children_of_tag[parent_tag]:
                        _children_of_tag[parent_tag].append(async_wrapper)
            return async_wrapper
        else:
            @functools.wraps(func)
            def sync_wrapper(*args, **kwargs):
                agent_id = f"{agent_name}_{uuid.uuid4().hex[:8]}"
                return _execute_sync(
                    func, args, kwargs, agent_name, agent_id, tag, notes, llm_input
                )
            sync_wrapper.__wrapped_tag__ = tag
            if tag:
                _tag_to_callable[tag] = sync_wrapper
                _tag_llm_input[tag] = llm_input
            if isinstance(llm_input, list) and len(llm_input) == 1:
                parent_tag = llm_input[0]
                parent_llm_input = _tag_llm_input.get(parent_tag, "UNKNOWN")
                if parent_llm_input is None:
                    _children_of_tag.setdefault(parent_tag, [])
                    if sync_wrapper not in _children_of_tag[parent_tag]:
                        _children_of_tag[parent_tag].append(sync_wrapper)
            return sync_wrapper


# ─────────────────────────────────────────────────────────────────────────────
# SYNC EXECUTION
# ─────────────────────────────────────────────────────────────────────────────

def _execute_sync(func, args, kwargs, agent_name, agent_id, tag, notes, llm_input):
    """
    Sync traced execution.
    parent_id is set ONLY when llm_input is a list of resolved tags.
    All other agents emit with parent_id=None (no edge in the graph).
    """
    run_id = _ensure_run()
    step = _next_step()
    input_data = _capture_input(func, args, kwargs)

    # Register tag so downstream agents can resolve it via llm_input
    if tag:
        with _state_lock:
            _tag_to_agent_ids.setdefault(tag, []).append(agent_id)

    # ONLY a list llm_input creates edges — string or None → no parent
    parent_id = _resolve_llm_input_parents(llm_input)

    start = time.monotonic()
    emit(build_event(
        event_type="agent_start",
        run_id=run_id,
        agent_id=agent_id,
        agent_name=agent_name,
        parent_id=parent_id,
        step_number=step,
        input=input_data,
        metadata={"tag": tag, "notes": notes, "llm_input": llm_input},
    ))

    try:
        result = func(*args, **kwargs)
        duration_ms = (time.monotonic() - start) * 1000
        emit(build_event(
            event_type="agent_end",
            run_id=run_id,
            agent_id=agent_id,
            agent_name=agent_name,
            parent_id=parent_id,
            step_number=step,
            input=input_data,
            output=result,
            tokens=extract_tokens(result, input_data),
            duration_ms=duration_ms,
            metadata={"tag": tag, "notes": notes, "llm_input": llm_input},
        ))
        return result
    except Exception as exc:
        duration_ms = (time.monotonic() - start) * 1000
        emit(build_event(
            event_type="agent_error",
            run_id=run_id,
            agent_id=agent_id,
            agent_name=agent_name,
            parent_id=parent_id,
            step_number=step,
            input=input_data,
            duration_ms=duration_ms,
            error={
                "message": str(exc),
                "type": type(exc).__name__,
                "traceback": traceback.format_exc(),
            },
            metadata={"tag": tag, "notes": notes, "llm_input": llm_input},
        ))
        raise


# ─────────────────────────────────────────────────────────────────────────────
# ASYNC EXECUTION
# ─────────────────────────────────────────────────────────────────────────────

async def _execute_async(func, args, kwargs, agent_name, agent_id, tag, notes, llm_input):
    """
    Async traced execution.
    parent_id is set ONLY when llm_input is a list of resolved tags.
    All other agents emit with parent_id=None (no edge in the graph).

    Auto-gather: when this agent has a registered `tag`, the library looks up
    all callables that declared llm_input=[this_tag] in _children_of_tag and
    spawns them in parallel automatically before running the function body.
    """
    run_id = _ensure_run()
    step = _next_step()
    input_data = _capture_input(func, args, kwargs)

    # Register tag FIRST so children that run during auto-gather can resolve it
    if tag:
        with _state_lock:
            _tag_to_agent_ids.setdefault(tag, []).append(agent_id)

    # parent_id resolved before emitting agent_start
    parent_id = _resolve_llm_input_parents(llm_input)

    start = time.monotonic()
    # Emit agent_start BEFORE spawning children so the node exists in the graph
    # when children emit their agent_start with parent_id pointing to this agent.
    emit(build_event(
        event_type="agent_start",
        run_id=run_id,
        agent_id=agent_id,
        agent_name=agent_name,
        parent_id=parent_id,
        step_number=step,
        input=input_data,
        metadata={"tag": tag, "notes": notes, "llm_input": llm_input},
    ))

    try:
        # ── Auto-gather children ──────────────────────────────────────────────
        # If any async callables registered themselves as children of this tag
        # (via llm_input=["this_tag"]), spawn them all in parallel now.
        # Only spawn children that haven't already run this session (their tag
        # not yet in _tag_to_agent_ids), so merge-point nodes don't re-spawn.
        auto_gathered = None
        if tag:
            children = _children_of_tag.get(tag, [])
            async_children = [c for c in children if asyncio.iscoroutinefunction(c)]
            if async_children:
                # Filter: skip children whose tag already has agent_ids registered
                to_spawn = []
                with _state_lock:
                    for c in async_children:
                        child_tag = getattr(c, "__wrapped_tag__", None)
                        if child_tag is None or child_tag not in _tag_to_agent_ids:
                            to_spawn.append(c)

                if to_spawn:
                    auto_gathered = list(await asyncio.gather(*[c(*args, **kwargs) for c in to_spawn]))

        if auto_gathered is not None:
            # Run the coordinator body (may do logging / post-processing).
            # If body returns None, return the gathered results list instead.
            result = await func(*args, **kwargs)
            if result is None:
                result = auto_gathered
        else:
            result = await func(*args, **kwargs)

        duration_ms = (time.monotonic() - start) * 1000
        emit(build_event(
            event_type="agent_end",
            run_id=run_id,
            agent_id=agent_id,
            agent_name=agent_name,
            parent_id=parent_id,
            step_number=step,
            input=input_data,
            output=result,
            tokens=extract_tokens(result, input_data),
            duration_ms=duration_ms,
            metadata={"tag": tag, "notes": notes, "llm_input": llm_input},
        ))
        return result
    except Exception as exc:
        duration_ms = (time.monotonic() - start) * 1000
        emit(build_event(
            event_type="agent_error",
            run_id=run_id,
            agent_id=agent_id,
            agent_name=agent_name,
            parent_id=parent_id,
            step_number=step,
            input=input_data,
            duration_ms=duration_ms,
            error={
                "message": str(exc),
                "type": type(exc).__name__,
                "traceback": traceback.format_exc(),
            },
            metadata={"tag": tag, "notes": notes, "llm_input": llm_input},
        ))
        raise


# ─────────────────────────────────────────────────────────────────────────────
# HELPERS
# ─────────────────────────────────────────────────────────────────────────────

def _capture_input(func: Callable, args: tuple, kwargs: dict) -> Any:
    """Extract meaningful input from function arguments."""
    try:
        sig = inspect.signature(func)
        params = list(sig.parameters.keys())
        bound = {}
        for i, arg in enumerate(args):
            if i < len(params):
                bound[params[i]] = arg
        bound.update(kwargs)
        if "self" in bound:
            del bound["self"]
        return bound if bound else None
    except Exception:
        return str(args[0]) if args else None
