"""Target factory functions for LangSmith experiments.

A target is a ``(inputs: dict) -> dict`` callable passed to ``ls_evaluate``.
Each factory builds and closes over the agent so it is constructed once and
shared across all dataset examples in a run.
"""

import asyncio
import json

import anthropic as anthropic_sdk
from langchain_core.language_models import BaseChatModel
from langchain_core.tools import BaseTool
from langsmith.wrappers import wrap_anthropic
from mcp import ClientSession
from mcp.client.sse import sse_client

from agents.multi import build_harmonic_graph
from langsmith_evals.fixtures import (
    _DYNAMIC_MOCKED_TOOLS,
    _SIMILARITY_POOL,
    _SIMILARITY_SCORE,
    load_eval_tools,
)

_SYSTEM = (
    "You are a harmonic analysis assistant. After analysing a chord sequence, present only "
    "well-known results — songs that are recognisable hits or by widely-known artists. "
    "Use your knowledge to filter out obscure tracks. If no well-known matches exist, say so."
)


def make_harmonic_target(
    mcp_url: str,
    analysis_llm: BaseChatModel,
    research_llm: BaseChatModel,
    search_tools: list[BaseTool],
) -> callable:
    """Target wrapping the full harmonic graph.

    Analysis similarity tools are mocked with fixture data so runs are
    deterministic and independent of the live FAISS index. Research uses
    the provided search tools, which may be live or mocked.

    :param mcp_url: SSE endpoint of the running MCP server.
    :param analysis_llm: LLM for the analysis subgraph.
    :param research_llm: LLM for the research subgraph.
    :param search_tools: Web search tools for the research subgraph.
    """
    mcp_tools = load_eval_tools(mcp_url)
    graph = build_harmonic_graph(analysis_llm, research_llm, mcp_tools, search_tools)

    def target(inputs: dict) -> dict:
        user_input = inputs["user_input"]
        print(f"\n>>> input: {user_input!r}")
        response = ""
        for ns, chunk in graph.stream({"user_input": user_input}, stream_mode="updates", subgraphs=True):
            for node_name, update in chunk.items():
                prefix = f"{ns[-1].split(':')[0]}/" if ns else ""
                print(f"  [node] {prefix}{node_name}")
                if not ns and "response" in update:
                    response = update["response"]
        print(f"<<< response: {response!r}")
        return {"response": response}

    return target


def make_claude_sdk_target(mcp_url: str, model: str) -> callable:
    """Target running the Anthropic SDK directly with native MCP.

    Uses ``wrap_anthropic`` so every API call is captured in LangSmith.
    Similarity tools are mocked with the same fixture pool used by the
    LangGraph experiments; all other tools call the live server.

    :param mcp_url: SSE endpoint of the running MCP server.
    :param model: Anthropic model ID.
    """
    sdk_client = wrap_anthropic(anthropic_sdk.Anthropic())

    async def _run(user_input: str) -> str:
        async with sse_client(mcp_url) as (read, write):
            async with ClientSession(read, write) as session:
                await session.initialize()
                tools = [
                    {"name": t.name, "description": t.description or "", "input_schema": t.inputSchema}
                    for t in (await session.list_tools()).tools
                ]
                messages = [{"role": "user", "content": user_input}]

                while True:
                    response = sdk_client.messages.create(
                        model=model,
                        max_tokens=4096,
                        system=_SYSTEM,
                        tools=tools,
                        messages=messages,
                    )

                    if response.stop_reason == "end_turn":
                        return next((b.text for b in response.content if b.type == "text"), "")

                    tool_results = []
                    for block in response.content:
                        if block.type != "tool_use":
                            continue
                        args = dict(block.input)
                        if block.name in _DYNAMIC_MOCKED_TOOLS:
                            limit = int(args.get("limit", len(_SIMILARITY_POOL)))
                            result_text = f"{_SIMILARITY_SCORE}\n{json.dumps(_SIMILARITY_POOL[:limit])}"
                        else:
                            mcp_result = await session.call_tool(block.name, args)
                            result_text = "\n".join(
                                item.text for item in mcp_result.content if hasattr(item, "text")
                            )
                        print(f"  [tool] {block.name}({args})")
                        tool_results.append({
                            "type": "tool_result",
                            "tool_use_id": block.id,
                            "content": result_text,
                        })

                    messages.append({"role": "assistant", "content": response.content})
                    messages.append({"role": "user", "content": tool_results})

    def target(inputs: dict) -> dict:
        user_input = inputs["user_input"]
        print(f"\n>>> input: {user_input!r}")
        response = asyncio.run(_run(user_input))
        print(f"<<< response: {response!r}")
        return {"response": response}

    return target