"""Client instrumentation for Claude Agent SDK."""

import logging
import time
import weakref
from collections.abc import AsyncGenerator, AsyncIterable
from datetime import datetime, timezone
from typing import Any, Optional

from langsmith._internal import _context
from langsmith._internal._package_version import get_package_version
from langsmith.run_helpers import get_current_run_tree, trace

from ._config import get_tracing_config
from ._hooks import (
    SessionState,
    _current_session,
    _register_session,
    _set_session_root,
    _unregister_session,
    clear_active_tool_runs,
    get_subagent_run_by_tool_id,
    post_tool_use_failure_hook,
    post_tool_use_hook,
    pre_tool_use_hook,
    subagent_start_hook,
    subagent_stop_hook,
)
from ._messages import (
    build_llm_input,
    flatten_content_blocks,
    unwrap_message_dicts,
)
from ._tools import (
    clear_parent_run_tree,
    get_parent_run_tree,
    set_parent_run_tree,
)
from ._transcripts import LLM_RUN_NAME, reconcile_from_transcripts
from ._usage import extract_usage_metadata

logger = logging.getLogger(__name__)

TRACE_CHAIN_NAME = "claude.conversation"


class TurnLifecycle:
    """Track ongoing model runs so consecutive messages are recorded correctly.

    The Claude Agent SDK may deliver a single assistant turn as multiple
    ``AssistantMessage`` events (e.g. one with ``ThinkingBlock``, another
    with ``TextBlock``/``ToolUseBlock``).  Messages that share the same
    ``message_id`` are accumulated into a single LLM run.
    """

    def __init__(self, query_start_time: Optional[float] = None):
        self.current_run: Optional[Any] = None
        self.current_message_id: Optional[str] = None
        self.next_start_time: Optional[float] = query_start_time
        # message_id → RunTree for all LLM runs created this conversation.
        # Used to retroactively set usage from transcripts.
        self.llm_runs_by_message_id: dict[str, Any] = {}
        # Runs that have been end()ed but not yet patch()ed.
        # Deferred so transcript usage can be set before the single patch().
        self._pending_patch: list[Any] = []

    def start_llm_run(
        self,
        message: Any,
        prompt: Any,
        history: list[dict[str, Any]],
        parent: Optional[Any] = None,
    ) -> Optional[dict[str, Any]]:
        """Begin or continue a model run for *message*.

        If *message* has the same ``message_id`` as the current run the
        output is appended; otherwise a new run is started (ending any
        previous one first).
        """
        message_id = getattr(message, "message_id", None)
        start = self.next_start_time or time.time()

        # Same turn – just accumulate the output blocks and update usage.
        # Return None so the caller does NOT append a duplicate history
        # entry; the original entry in ``history`` is updated in place.
        if message_id and message_id == self.current_message_id and self.current_run:
            content = flatten_content_blocks(getattr(message, "content", None))
            if content and self.current_run.outputs:
                prev = self.current_run.outputs.get("content", [])
                if isinstance(prev, list) and isinstance(content, list):
                    merged = prev + content
                    self.current_run.outputs["content"] = merged
                    # Update the existing history entry in place so
                    # subsequent LLM runs see a single merged message.
                    for entry in reversed(history):
                        if entry.get("role") == "assistant":
                            entry["content"] = merged
                            break
                elif isinstance(content, list):
                    self.current_run.outputs["content"] = content
            self._set_usage_from_message(message, self.current_run)
            return None

        # Different turn – end previous but defer patch() until
        # transcript usage is available.
        if self.current_run:
            self.current_run.end()
            self._pending_patch.append(self.current_run)

        final_output, run = begin_llm_run_from_assistant_messages(
            [message], prompt, history, start_time=start, parent=parent
        )
        self.current_run = run
        self.current_message_id = message_id
        self.next_start_time = None

        if run:
            if message_id:
                self.llm_runs_by_message_id[message_id] = run
            self._set_usage_from_message(message, run)

        return final_output

    @staticmethod
    def _set_usage_from_message(message: Any, run: Any) -> None:
        """Set usage metadata on a run from a live AssistantMessage.

        Always overwrites — later chunks in the same turn have more
        accurate counts.  Transcript-based usage will overwrite again
        if available.
        """
        raw_usage = getattr(message, "usage", None)
        if not raw_usage:
            return
        usage_meta = extract_usage_metadata(raw_usage)
        if usage_meta:
            meta = run.extra.setdefault("metadata", {})
            meta["usage_metadata"] = usage_meta

    def mark_next_start(self) -> None:
        """Mark when the next assistant message will start."""
        self.next_start_time = time.time()

    def close(self) -> None:
        """End any open run and add to pending patch list."""
        if self.current_run:
            self.current_run.end()
            self._pending_patch.append(self.current_run)
            self.current_run = None

    def flush(self) -> None:
        """Patch all deferred LLM runs. Call after usage has been set."""
        for run in self._pending_patch:
            try:
                run.patch()
            except Exception as e:
                logger.warning(f"Failed to patch LLM run: {e}")
        self._pending_patch.clear()


def begin_llm_run_from_assistant_messages(
    messages: list[Any],
    prompt: Any,
    history: list[dict[str, Any]],
    start_time: Optional[float] = None,
    parent: Optional[Any] = None,
) -> tuple[Optional[dict[str, Any]], Optional[Any]]:
    """Create a traced model run from assistant messages."""
    if not messages or type(messages[-1]).__name__ != "AssistantMessage":
        return None, None

    last_msg = messages[-1]
    model = getattr(last_msg, "model", None)
    if parent is None:
        parent = get_parent_run_tree() or get_current_run_tree()
    if not parent:
        return None, None

    inputs = build_llm_input(prompt, history)
    outputs = [
        {"content": flatten_content_blocks(m.content), "role": "assistant"}
        for m in messages
        if hasattr(m, "content")
    ]

    llm_metadata: dict[str, Any] = {"ls_provider": "anthropic"}
    if model:
        llm_metadata["ls_model_name"] = model

    llm_run = parent.create_child(
        name=LLM_RUN_NAME,
        run_type="llm",
        inputs={"messages": inputs} if inputs else {},
        extra={"metadata": llm_metadata},
        start_time=datetime.fromtimestamp(start_time, tz=timezone.utc)
        if start_time
        else None,
    )

    try:
        llm_run.post()
    except Exception as e:
        logger.warning(f"Failed to post LLM run: {e}")

    # Set outputs after posting so they are sent with end_time on the patch.
    llm_run.outputs = outputs[-1] if len(outputs) == 1 else {"content": outputs}

    final_content = (
        {"content": flatten_content_blocks(last_msg.content), "role": "assistant"}
        if hasattr(last_msg, "content")
        else None
    )
    return final_content, llm_run


def _bind_hook_to_session(hook: Any, session: Optional[SessionState]) -> Any:
    """Return a hook callable that runs with *session* bound, if provided."""
    if session is None:
        return hook

    async def _bound(input_data: Any, tool_use_id: Any, context: Any) -> Any:
        token = _current_session.set(session)
        try:
            return await hook(input_data, tool_use_id, context)
        finally:
            _current_session.reset(token)

    return _bound


def _inject_tracing_hooks(options: Any, session: Optional[SessionState] = None) -> None:
    """Inject LangSmith tracing hooks into ClaudeAgentOptions.

    If *session* is provided, injected hook callables bind that session around
    each hook invocation. This is important because the Claude SDK may execute
    hooks in async contexts that do not inherit the ``receive_response``
    ContextVar; binding at hook injection time keeps each client isolated.
    """
    if not hasattr(options, "hooks"):
        return

    # Initialize hooks dict if not present
    if options.hooks is None:
        options.hooks = {}

    for event in (
        "PreToolUse",
        "PostToolUse",
        "PostToolUseFailure",
        "SubagentStart",
        "SubagentStop",
    ):
        if event not in options.hooks:
            options.hooks[event] = []

    try:
        from claude_agent_sdk import HookMatcher  # type: ignore[import-not-found]

        langsmith_pre_matcher = HookMatcher(
            matcher=None, hooks=[_bind_hook_to_session(pre_tool_use_hook, session)]
        )
        langsmith_post_matcher = HookMatcher(
            matcher=None, hooks=[_bind_hook_to_session(post_tool_use_hook, session)]
        )
        langsmith_failure_matcher = HookMatcher(
            matcher=None,
            hooks=[_bind_hook_to_session(post_tool_use_failure_hook, session)],
        )
        langsmith_subagent_start_matcher = HookMatcher(
            matcher=None, hooks=[_bind_hook_to_session(subagent_start_hook, session)]
        )
        langsmith_subagent_stop_matcher = HookMatcher(
            matcher=None, hooks=[_bind_hook_to_session(subagent_stop_hook, session)]
        )

        options.hooks["PreToolUse"].insert(0, langsmith_pre_matcher)
        options.hooks["PostToolUse"].insert(0, langsmith_post_matcher)
        options.hooks["PostToolUseFailure"].insert(0, langsmith_failure_matcher)
        options.hooks["SubagentStart"].insert(0, langsmith_subagent_start_matcher)
        options.hooks["SubagentStop"].insert(0, langsmith_subagent_stop_matcher)

        logger.debug("Injected LangSmith tracing hooks into ClaudeAgentOptions")
    except ImportError:
        logger.warning("Failed to import HookMatcher from claude_agent_sdk")
    except Exception as e:
        logger.warning(f"Failed to inject tracing hooks: {e}")


def _wrap_tool_handler(
    original_handler: Any,
    session: Optional[SessionState] = None,
    tool_name: Optional[str] = None,
) -> Any:
    """Wrap an MCP tool handler to propagate LangSmith run context.

    The Claude SDK runs hooks and tool handlers in different async task
    contexts, so contextvars set in ``PreToolUse`` are invisible to the
    handler. This wrapper copies the active tool run into the contextvar before
    calling the original handler, so ``@traceable`` calls inside the handler
    nest correctly.
    """

    async def _wrapped(args: Any) -> Any:
        # The most recently added active tool run is the one PreToolUse just
        # created for this invocation. Prefer an explicitly bound client
        # session because tool handlers may run in an async context that did
        # not inherit _current_session.
        tool_run = _get_last_active_tool_run(session, args=args, tool_name=tool_name)
        if tool_run:
            token = _context._PARENT_RUN_TREE_REF.set(weakref.ref(tool_run))
            session_token = (
                _current_session.set(session) if session is not None else None
            )
            try:
                return await original_handler(args)
            finally:
                if session_token is not None:
                    _current_session.reset(session_token)
                _context._PARENT_RUN_TREE_REF.reset(token)
        return await original_handler(args)

    _wrapped._langsmith_wrapped = True  # type: ignore[attr-defined]
    _wrapped._langsmith_original_handler = original_handler  # type: ignore[attr-defined]
    _wrapped._langsmith_session = session  # type: ignore[attr-defined]
    _wrapped._langsmith_tool_name = tool_name  # type: ignore[attr-defined]
    return _wrapped


def _tool_run_matches(run: Any, args: Any, tool_name: Optional[str]) -> bool:
    """Return whether *run* appears to be for this SDK MCP handler call.

    Matching is intentionally strict: we require both the tool name and the
    handler args to line up with what the ``PreToolUse`` hook recorded. This
    avoids cross-attributing a handler invocation to the wrong client's active
    tool run under concurrency.
    """
    if not tool_name:
        return False
    run_name = str(getattr(run, "name", ""))
    # SDK MCP tools show up in hook data as e.g. ``mcp__weather__get_weather``
    # while the handler only knows its short name ``get_weather``.
    name_matches = (
        tool_name == run_name or tool_name in run_name or run_name in tool_name
    )
    if not name_matches:
        return False
    inputs = getattr(run, "inputs", None)
    if not isinstance(inputs, dict):
        return False
    # PreToolUse stores {} when the tool had no inputs, otherwise
    # {"input": <tool_input>}. Normalise both sides before comparing.
    recorded = inputs.get("input", {}) if inputs else {}
    return recorded == (args or {})


def _newest_matching_tool_run(
    sessions: list[SessionState], args: Any, tool_name: Optional[str]
) -> Any:
    """Return the most recently created active tool run that matches."""
    candidates: list[tuple[float, Any]] = []
    for candidate_session in sessions:
        for run, start_time in candidate_session.active_tool_runs.values():
            if _tool_run_matches(run, args, tool_name):
                candidates.append((start_time, run))
    if not candidates:
        return None
    return max(candidates, key=lambda item: item[0])[1]


def _get_last_active_tool_run(
    session: Optional[SessionState] = None,
    *,
    args: Any = None,
    tool_name: Optional[str] = None,
) -> Any:
    """Return the active tool run for an SDK MCP handler, or None.

    Lookup order:

    1. The session explicitly bound to the handler (if any).
    2. The session bound to the current ContextVar.
    3. The module-level default session (unit tests / unbound callers).
    4. Any live client session — only used when the handler is unbound and the
       SDK invoked it in a detached async context. Requires an exact tool
       name + args match to avoid cross-attribution across clients.
    """
    from ._hooks import (
        _current_session,
        _current_session_or_default,
        _registered_sessions,
    )

    # If we have a specific session (explicitly bound, current-context, or the
    # test default), just return its newest active tool run. There is no
    # cross-client ambiguity at that point.
    def _newest_in(s: SessionState) -> Any:
        if not s.active_tool_runs:
            return None
        latest_id = max(
            s.active_tool_runs,
            key=lambda tid: s.active_tool_runs[tid][1],
        )
        return s.active_tool_runs[latest_id][0]

    if session is not None:
        return _newest_in(session)

    current_session = _current_session.get()
    if current_session is not None:
        run = _newest_in(current_session)
        if run is not None:
            return run

    default_session = _current_session_or_default()
    if default_session is not current_session:
        run = _newest_in(default_session)
        if run is not None:
            return run

    # Last resort: the SDK invoked this handler in a detached async context and
    # the handler object wasn't bound to a session. Require strict tool-name +
    # args match so concurrent clients can't steal each other's attribution.
    return _newest_matching_tool_run(_registered_sessions(), args, tool_name)


def instrument_claude_client(original_class: Any) -> None:
    """Patch ``ClaudeSDKClient`` **in place** to trace calls.

    In-place patching (rather than subclassing + reference replacement)
    ensures that callers who imported ``ClaudeSDKClient`` *before*
    ``configure_claude_agent_sdk()`` was called still get instrumented.
    """
    if getattr(original_class, "_langsmith_instrumented", False):
        return  # Already wrapped, avoid double-tracing

    # ── stash originals ──────────────────────────────────────────────
    _orig_init = original_class.__init__
    _orig_query = original_class.query
    _orig_receive_response = original_class.receive_response

    # ── patched __init__ ─────────────────────────────────────────────
    def _traced_init(self: Any, *args: Any, **kwargs: Any) -> None:
        options = kwargs.get("options") or (args[0] if args else None)
        self._ls_session = SessionState()
        if options:
            _inject_tracing_hooks(options, self._ls_session)
        _orig_init(self, *args, **kwargs)
        self._ls_prompt = None
        self._ls_start_time = None
        self._ls_streamed_input = None

    # ── patched query ────────────────────────────────────────────────
    async def _traced_query(self: Any, *args: Any, **kwargs: Any) -> Any:
        self._ls_start_time = time.time()
        self._ls_streamed_input = None
        prompt = args[0] if args else kwargs.get("prompt")

        if prompt is None:
            pass
        elif isinstance(prompt, str):
            self._ls_prompt = prompt
        elif isinstance(prompt, AsyncIterable):
            collector: list[dict[str, Any]] = []
            self._ls_streamed_input = collector
            self._ls_prompt = None

            async def _gen_wrapper() -> AsyncGenerator[dict[str, Any], None]:
                async for msg in prompt:
                    collector.append(msg)
                    yield msg

            if args:
                args = (_gen_wrapper(),) + args[1:]
            else:
                kwargs["prompt"] = _gen_wrapper()
        else:
            self._ls_prompt = str(prompt)

        return await _orig_query(self, *args, **kwargs)

    # ── patched receive_response ─────────────────────────────────────
    async def _traced_receive_response(self: Any) -> AsyncGenerator[Any, None]:
        messages = _orig_receive_response(self)

        trace_inputs: dict[str, Any] = {}
        trace_metadata: dict[str, Any] = {
            "ls_integration": "claude-agent-sdk",
            "ls_integration_version": get_package_version("claude_agent_sdk"),
        }

        awaiting_streamed_input = self._ls_streamed_input is not None

        if self._ls_prompt:
            trace_inputs["prompt"] = self._ls_prompt

        if hasattr(self, "options") and self.options:
            if hasattr(self.options, "system_prompt") and self.options.system_prompt:
                system_prompt = self.options.system_prompt
                if isinstance(system_prompt, str):
                    trace_inputs["system"] = system_prompt
                elif isinstance(system_prompt, dict):
                    if system_prompt.get("type") == "preset":
                        preset_text = (
                            f"preset: {system_prompt.get('preset', 'claude_code')}"
                        )
                        if "append" in system_prompt:
                            preset_text += f"\nappend: {system_prompt['append']}"
                        trace_inputs["system"] = preset_text
                    else:
                        trace_inputs["system"] = system_prompt

            for attr in ["model", "permission_mode", "max_turns"]:
                if hasattr(self.options, attr):
                    val = getattr(self.options, attr)
                    if val is not None:
                        trace_metadata[attr] = val

        config = get_tracing_config()
        user_metadata = config.get("metadata") or {}

        trace_kwargs: dict[str, Any] = {
            "name": config.get("name") or TRACE_CHAIN_NAME,
            "run_type": "chain",
            "inputs": trace_inputs,
            "metadata": {
                **trace_metadata,
                **user_metadata,
                "ls_agent_type": "root",
            },
        }
        if config.get("project_name"):
            trace_kwargs["project_name"] = config["project_name"]
        if config.get("tags"):
            trace_kwargs["tags"] = config["tags"]

        async with trace(**trace_kwargs) as run:
            # Bind this client's state container to the ContextVar so stream
            # helpers on this SDK event loop pick it up (see
            # _hooks.SessionState). This keeps concurrent ClaudeSDKClient
            # instances — eval runs, FastAPI handlers, Celery workers,
            # asyncio.gather — from corrupting each other's correlation state.
            session = getattr(self, "_ls_session", None)
            if session is None:
                session = self._ls_session = SessionState()
            session_token = _register_session(session)
            _set_session_root(session, run)
            parent_token = set_parent_run_tree(run)
            tracker = TurnLifecycle(self._ls_start_time)
            collected_by_ctx: dict[Optional[str], list[dict[str, Any]]] = {None: []}

            prompt_for_llm: Any = self._ls_prompt

            try:
                async for msg in messages:
                    if awaiting_streamed_input and self._ls_streamed_input:
                        unwrapped_messages = unwrap_message_dicts(
                            self._ls_streamed_input
                        )
                        if unwrapped_messages:
                            run.inputs["messages"] = unwrapped_messages
                            prompt_for_llm = self._ls_streamed_input
                        awaiting_streamed_input = False

                    msg_type = type(msg).__name__

                    if msg_type == "AssistantMessage":
                        parent_tool_use_id = getattr(msg, "parent_tool_use_id", None)
                        llm_parent = (
                            get_subagent_run_by_tool_id(parent_tool_use_id)
                            if parent_tool_use_id
                            else None
                        )

                        ctx_key = parent_tool_use_id
                        ctx_history = collected_by_ctx.setdefault(ctx_key, [])

                        content = tracker.start_llm_run(
                            msg,
                            prompt_for_llm if parent_tool_use_id is None else None,
                            ctx_history,
                            parent=llm_parent,
                        )
                        if content:
                            ctx_history.append(content)

                    elif msg_type == "UserMessage":
                        parent_tool_use_id = getattr(msg, "parent_tool_use_id", None)
                        ctx_key = parent_tool_use_id
                        ctx_history = collected_by_ctx.setdefault(ctx_key, [])

                        if hasattr(msg, "content"):
                            flattened = flatten_content_blocks(msg.content)
                            if (
                                isinstance(flattened, list)
                                and flattened
                                and isinstance(flattened[0], dict)
                                and flattened[0].get("type") == "tool_result"
                            ):
                                for block in flattened:
                                    tool_use_id = block.get("tool_use_id")
                                    ctx_history.append(
                                        {
                                            "role": "tool",
                                            "content": block.get("content", ""),
                                            "tool_call_id": tool_use_id,
                                        }
                                    )
                                    if (
                                        tool_use_id
                                        and tool_use_id in session.active_tool_runs
                                    ):
                                        tool_run, _ = session.active_tool_runs.pop(
                                            tool_use_id
                                        )
                                        result_content = block.get("content", "")
                                        is_error = block.get("is_error", False)
                                        tool_run.end(
                                            outputs={"output": result_content},
                                            error=str(result_content)
                                            if is_error
                                            else None,
                                        )
                                        try:
                                            tool_run.patch()
                                        except Exception as e:
                                            logger.warning(
                                                "Failed to patch"
                                                f" orphaned tool run: {e}"
                                            )
                            else:
                                ctx_history.append(
                                    {
                                        "content": flattened,
                                        "role": "user",
                                    }
                                )
                        tracker.mark_next_start()
                    elif msg_type == "ResultMessage":
                        session_id_val = getattr(msg, "session_id", None)
                        meta = {
                            k: v
                            for k, v in {
                                "num_turns": getattr(msg, "num_turns", None),
                                "session_id": session_id_val,
                                "thread_id": session_id_val,
                                "duration_ms": getattr(msg, "duration_ms", None),
                                "duration_api_ms": getattr(
                                    msg, "duration_api_ms", None
                                ),
                                "is_error": getattr(msg, "is_error", None),
                            }.items()
                            if v is not None
                        }
                        if meta:
                            run.metadata.update(meta)

                    yield msg
                main_collected = collected_by_ctx.get(None, [])
                run.end(outputs=main_collected[-1] if main_collected else None)
            except Exception:
                logger.exception("Error while tracing Claude Agent stream")
            finally:
                tracker.close()
                reconcile_from_transcripts(tracker, session=session)
                tracker.flush()
                clear_parent_run_tree(parent_token)
                try:
                    clear_active_tool_runs(session)
                finally:
                    _unregister_session(session, session_token)

    # ── apply patches to the class itself ────────────────────────────
    original_class.__init__ = _traced_init
    original_class.query = _traced_query
    original_class.receive_response = _traced_receive_response
    original_class._langsmith_instrumented = True


def instrument_sdk_mcp_tool(tool_class: Any) -> None:
    """Patch ``SdkMcpTool.__init__`` to auto-wrap handlers.

    Wrapping happens at construction time so that any tool created
    *after* ``configure_claude_agent_sdk()`` automatically gets
    run-context propagation, regardless of how ``tool`` or
    ``create_sdk_mcp_server`` were imported.
    """
    if getattr(tool_class, "_langsmith_handler_patched", False):
        return

    _orig_init = tool_class.__init__

    def _patched_init(self: Any, *args: Any, **kwargs: Any) -> None:
        _orig_init(self, *args, **kwargs)
        handler = self.handler
        if callable(handler) and not getattr(handler, "_langsmith_wrapped", False):
            self.handler = _wrap_tool_handler(
                handler, tool_name=getattr(self, "name", None)
            )

    tool_class.__init__ = _patched_init
    tool_class._langsmith_handler_patched = True
