"""
aaiclick.ai.agents.debug_agent - Tier 1 lineage debugger.

`debug_result()` runs an LLM tool loop scoped to the backward lineage graph of
the target table. The agent reads rendered SQL templates, queries live tables
via ``LineageToolbox``, and produces an explanation citing concrete evidence.

See ``docs/designs/lineage.md``.
"""

from __future__ import annotations

import asyncio
import json
import logging
from typing import Any, NamedTuple

from aaiclick.ai.agents.lineage_tools import LINEAGE_TOOL_DEFINITIONS, LineageToolbox
from aaiclick.ai.agents.prompts import LINEAGE_TIER1_SYSTEM_PROMPT
from aaiclick.ai.config import get_ai_provider
from aaiclick.oplog.lineage import OplogGraph, oplog_subgraph

logger = logging.getLogger(__name__)


class _ToolCallSpec(NamedTuple):
    id: str
    name: str
    arguments: str


async def debug_result(
    target_table: str,
    question: str,
    graph: OplogGraph | None = None,
    max_iterations: int = 10,
) -> str:
    """Tier 1 lineage debug loop.

    The agent starts with the question, the backward oplog graph of
    `target_table`, and the four graph-scoped tools from ``LineageToolbox``.
    It iterates up to `max_iterations` rounds of tool calls before producing
    a final explanation.

    Pass `graph` to reuse a pre-built backward lineage graph and skip
    traversal.
    """
    if graph is None:
        graph = await oplog_subgraph(target_table, direction="backward")
    labels = graph.build_labels()
    toolbox = LineageToolbox(graph)

    context = graph.to_prompt_context()
    nodes = await toolbox.list_graph_nodes()
    liveness_lines = [f"- {n.table} [{n.kind}] live={n.live}" for n in nodes]
    context += "\n\n# Graph Node Liveness\n" + "\n".join(liveness_lines)

    user_content = f"Context:\n{context}\n\nTarget table: `{target_table}`\n\nQuestion: {question}"

    messages: list[dict[str, Any]] = [
        {"role": "system", "content": LINEAGE_TIER1_SYSTEM_PROMPT},
        {"role": "user", "content": user_content},
    ]

    provider = get_ai_provider()
    try:
        for _ in range(max_iterations):
            response = await provider.complete(messages, tools=LINEAGE_TOOL_DEFINITIONS)
            choice = response.choices[0]
            message = choice.message

            if message.tool_calls:
                calls = [_ToolCallSpec(tc.id, tc.function.name, tc.function.arguments) for tc in message.tool_calls]
                assistant_content = message.content
            else:
                # Ollama-served models (qwen2.5, llama3.1, ...) sometimes emit a tool
                # call as a JSON blob in `content` instead of the structured
                # `tool_calls` field. Salvage the call so the loop continues
                # instead of returning the JSON to the user as a final answer.
                inline = _extract_inline_tool_calls(message.content)
                if not inline:
                    return OplogGraph.replace_labels(message.content or "", labels)
                calls = inline
                assistant_content = ""

            messages.append(
                {
                    "role": "assistant",
                    "content": assistant_content,
                    "tool_calls": [
                        {
                            "id": calls[0].id,
                            "type": "function",
                            "function": {"name": calls[0].name, "arguments": calls[0].arguments},
                        }
                    ],
                }
            )

            tool_results = await asyncio.gather(
                *(toolbox.dispatch_tool(call.name, _parse_arguments(call.arguments)) for call in calls)
            )
            for index, (call, result) in enumerate(zip(calls, tool_results, strict=False)):
                messages.append(
                    {
                        "role": "tool",
                        "tool_call_id": call.id,
                        "content": result,
                    }
                )
                if index + 1 < len(calls):
                    next_call = calls[index + 1]
                    messages.append(
                        {
                            "role": "assistant",
                            "content": "",
                            "tool_calls": [
                                {
                                    "id": next_call.id,
                                    "type": "function",
                                    "function": {"name": next_call.name, "arguments": next_call.arguments},
                                }
                            ],
                        }
                    )

        messages.append({"role": "user", "content": "Please provide your final answer."})
        response = await provider.complete(messages)
        return OplogGraph.replace_labels(response.choices[0].message.content or "", labels)
    except Exception as exc:
        logger.warning("debug_result provider call failed: %s", exc)
        return f"(debug agent did not converge: {exc})"


def _parse_arguments(raw: str | None) -> dict[str, Any]:
    """Parse a tool call's JSON arguments, returning {} on malformed input.

    A model occasionally emits empty-string arguments for zero-arg tools;
    tolerate that instead of aborting the loop.
    """
    if not raw:
        return {}
    try:
        parsed = json.loads(raw)
    except json.JSONDecodeError:
        logger.warning("tool arguments not valid JSON: %r", raw)
        return {}
    return parsed if isinstance(parsed, dict) else {}


def _extract_inline_tool_calls(content: str | None) -> list[_ToolCallSpec]:
    """Recover tool calls that a model emitted as JSON text in `content`.

    Returns ``[]`` when ``content`` is empty, isn't JSON, or doesn't match the
    ``{"function": {"name": ..., "arguments": ...}}`` shape — the caller treats
    that as a true final answer.
    """
    if not content:
        return []
    text = content.strip()
    if not text or text[0] not in "{[":
        return []
    try:
        parsed = json.loads(text)
    except json.JSONDecodeError:
        return []
    items = parsed if isinstance(parsed, list) else [parsed]
    calls: list[_ToolCallSpec] = []
    for index, item in enumerate(items):
        if not isinstance(item, dict):
            return []
        fn = item.get("function")
        if not isinstance(fn, dict):
            return []
        name = fn.get("name")
        if not isinstance(name, str):
            return []
        raw_args = fn.get("arguments", {})
        if isinstance(raw_args, dict):
            arguments = json.dumps(raw_args)
        elif isinstance(raw_args, str):
            arguments = raw_args
        else:
            return []
        call_id = item.get("id")
        if not isinstance(call_id, str) or not call_id:
            call_id = f"inline_call_{index}"
        calls.append(_ToolCallSpec(id=call_id, name=name, arguments=arguments))
    return calls
