"""
Schema Adapters — normalize non-ReAct agent traces to QPPG chain format.

QPPG was built for HotpotQA/TriviaQA ReAct-style agents (Thought / Action:
Search[...] / Observation / Action: Finish[...]). Modern LLM APIs produce
very different trace formats. This module normalizes them all to the shared
QPPG step schema so LabelFreeScorer works out of the box.

Supported formats
-----------------
  1. ReAct (native) — no transformation needed
  2. OpenAI function-calling — chat completions with tool_calls / tool results
  3. OpenAI code-interpreter — assistant writes + executes Python; outputs results
  4. Structured output (tool-use) — Anthropic / OpenAI structured JSON responses
  5. LangChain AgentAction log — pre-v0.2 dict format with AgentAction records
  6. AutoDetect — inspect the chain dict and pick the right adapter

QPPG step schema (target format)
---------------------------------
    {
        "thought":     str,   # reasoning before the action
        "action_type": str,   # "Search", "Code", "Tool", "Finish"
        "action_arg":  str,   # argument/input to the action
        "observation": str,   # result/output from the action
    }

Usage
-----
    from qppg_service.integrations.schema_adapters import adapt_chain

    # Works on any of the 5 formats above
    qppg_chain = adapt_chain(raw_chain)
    # qppg_chain["steps"] is now in QPPG format
    result = scorer.score(qppg_chain["question"], qppg_chain["steps"],
                          qppg_chain["final_answer"], qppg_chain["finished"])
"""

from __future__ import annotations

import json
import re
from typing import Any


# ── Shared output schema ──────────────────────────────────────────────────────

def _step(thought: str = "", action_type: str = "Tool",
          action_arg: str = "", observation: str = "") -> dict:
    return {
        "thought":     (thought     or "")[:2000],
        "action_type": (action_type or "Tool")[:64],
        "action_arg":  (action_arg  or "")[:2000],
        "observation": (observation or "")[:2000],
    }


# ── 1. ReAct (native) ─────────────────────────────────────────────────────────

def adapt_react(chain: dict) -> dict:
    """
    Pass-through for native QPPG/ReAct chains.

    Expected input keys: question, steps, final_answer, finished
    Each step already has: thought, action_type, action_arg, observation
    """
    return {
        "question":     chain.get("question", ""),
        "steps":        chain.get("steps", []),
        "final_answer": chain.get("final_answer", ""),
        "finished":     bool(chain.get("finished", True)),
    }


# ── 2. OpenAI function-calling ────────────────────────────────────────────────

def adapt_openai_function_calls(messages: list[dict], question: str = "") -> dict:
    """
    Normalize an OpenAI chat-completion message history with tool_calls.

    Input format (messages list):
        [
            {"role": "user", "content": "What is X?"},
            {"role": "assistant", "content": null,
             "tool_calls": [{"id": "...", "function": {"name": "search",
                             "arguments": '{"query": "X"}'}}]},
            {"role": "tool", "tool_call_id": "...", "content": "...result..."},
            {"role": "assistant", "content": "The answer is ..."},
        ]
    """
    if not question:
        for m in messages:
            if m.get("role") == "user" and m.get("content"):
                question = str(m["content"])[:500]
                break

    steps        = []
    final_answer = ""
    finished     = False

    i = 0
    while i < len(messages):
        msg  = messages[i]
        role = msg.get("role", "")

        if role == "assistant":
            content    = msg.get("content") or ""
            tool_calls = msg.get("tool_calls") or []

            if tool_calls:
                # One tool call per step (take the first if batched)
                tc          = tool_calls[0]
                fn          = tc.get("function", {})
                fn_name     = fn.get("name", "unknown")
                fn_args_raw = fn.get("arguments", "{}")
                try:
                    fn_args = json.loads(fn_args_raw)
                    fn_arg_str = next(iter(fn_args.values()), fn_args_raw) if fn_args else fn_args_raw
                except Exception:
                    fn_arg_str = fn_args_raw

                # Peek ahead for the tool result
                observation = ""
                if i + 1 < len(messages) and messages[i + 1].get("role") == "tool":
                    observation = str(messages[i + 1].get("content", ""))
                    i += 1  # consume the tool message

                action_type = _normalise_tool_name(fn_name)
                steps.append(_step(
                    thought     = content,
                    action_type = action_type,
                    action_arg  = str(fn_arg_str),
                    observation = observation,
                ))
            elif content:
                # Final assistant message — the answer
                final_answer = content
                finished     = True
                steps.append(_step(
                    thought     = "",
                    action_type = "Finish",
                    action_arg  = content,
                    observation = "",
                ))
        i += 1

    return {
        "question":     question,
        "steps":        steps,
        "final_answer": final_answer,
        "finished":     finished,
    }


# ── 3. OpenAI code-interpreter ────────────────────────────────────────────────

def adapt_code_interpreter(messages: list[dict], question: str = "") -> dict:
    """
    Normalize an OpenAI Assistants API code-interpreter run.

    Input format:
        [
            {"role": "user", "content": "Compute X"},
            {"role": "assistant", "content": [
                {"type": "text", "text": {"value": "Let me compute..."}},
                {"type": "code", "text": {"value": "import math; print(math.sqrt(2))"}}
             ]},
            {"role": "tool", "content": [
                {"type": "logs", "logs": "1.4142135623730951"}
             ]},
            {"role": "assistant", "content": [
                {"type": "text", "text": {"value": "The answer is 1.41"}}
             ]},
        ]
    Also handles flat-string content for simpler formats.
    """
    if not question:
        for m in messages:
            if m.get("role") == "user":
                c = m.get("content", "")
                question = (c if isinstance(c, str) else str(c))[:500]
                break

    steps        = []
    final_answer = ""
    finished     = False

    def _extract_text(content) -> str:
        if isinstance(content, str):
            return content
        if isinstance(content, list):
            parts = []
            for block in content:
                if isinstance(block, dict):
                    if block.get("type") == "text":
                        parts.append(block.get("text", {}).get("value", ""))
                    elif block.get("type") in ("code", "input"):
                        parts.append(block.get("text", {}).get("value", ""))
                    elif block.get("type") == "logs":
                        parts.append(block.get("logs", ""))
            return "\n".join(parts)
        return str(content)

    i = 0
    while i < len(messages):
        msg  = messages[i]
        role = msg.get("role", "")

        if role == "assistant":
            content_raw = msg.get("content", "")
            content     = _extract_text(content_raw)

            # Check if next message is a tool/code output
            observation = ""
            if i + 1 < len(messages) and messages[i + 1].get("role") in ("tool", "code_result"):
                observation = _extract_text(messages[i + 1].get("content", ""))
                i += 1

            if observation:
                steps.append(_step(
                    thought     = content,
                    action_type = "Code",
                    action_arg  = content,
                    observation = observation,
                ))
            else:
                final_answer = content
                finished     = True
                steps.append(_step(
                    thought     = "",
                    action_type = "Finish",
                    action_arg  = content,
                ))
        i += 1

    return {
        "question":     question,
        "steps":        steps,
        "final_answer": final_answer,
        "finished":     finished,
    }


# ── 4. Structured / tool-use output ───────────────────────────────────────────

def adapt_structured_output(trace: dict) -> dict:
    """
    Normalize a structured tool-use trace (Anthropic / OpenAI tool-use).

    Input format:
        {
            "question": "...",
            "turns": [
                {"role": "assistant", "content": [
                    {"type": "thinking", "thinking": "..."},
                    {"type": "tool_use", "name": "web_search",
                     "input": {"query": "..."}},
                ]},
                {"role": "user", "content": [
                    {"type": "tool_result", "content": "...result..."}
                ]},
                {"role": "assistant", "content": [
                    {"type": "text", "text": "The answer is..."}
                ]}
            ]
        }
    """
    question     = trace.get("question", "")
    turns        = trace.get("turns", [])
    steps        = []
    final_answer = ""
    finished     = False

    i = 0
    while i < len(turns):
        turn = turns[i]
        role = turn.get("role", "")

        if role == "assistant":
            content = turn.get("content", [])
            if isinstance(content, str):
                content = [{"type": "text", "text": content}]

            thought    = ""
            tool_steps = []

            for block in content:
                btype = block.get("type", "")
                if btype in ("thinking", "reasoning"):
                    thought = block.get("thinking") or block.get("text", "")
                elif btype == "tool_use":
                    name      = block.get("name", "unknown")
                    inp       = block.get("input", {})
                    arg_str   = next(iter(inp.values()), json.dumps(inp)) if inp else ""
                    tool_steps.append((name, str(arg_str), thought))
                    thought   = ""  # reset after consuming
                elif btype == "text":
                    text = block.get("text", "")
                    if text and not tool_steps:
                        thought = text

            # Peek ahead for tool results
            obs_blocks = []
            if i + 1 < len(turns) and turns[i + 1].get("role") == "user":
                next_content = turns[i + 1].get("content", [])
                if isinstance(next_content, list):
                    for b in next_content:
                        if b.get("type") == "tool_result":
                            obs_blocks.append(str(b.get("content", "")))
                if obs_blocks:
                    i += 1  # consume the user/tool_result turn

            if tool_steps:
                for j, (name, arg, tht) in enumerate(tool_steps):
                    obs = obs_blocks[j] if j < len(obs_blocks) else ""
                    steps.append(_step(
                        thought     = tht,
                        action_type = _normalise_tool_name(name),
                        action_arg  = arg,
                        observation = obs,
                    ))
            elif content:
                # Pure text response — final answer
                text = " ".join(
                    b.get("text", "") for b in (content if isinstance(content, list) else [])
                    if isinstance(b, dict) and b.get("type") == "text"
                ) or (content if isinstance(content, str) else "")
                final_answer = text
                finished     = True
                steps.append(_step(
                    thought     = "",
                    action_type = "Finish",
                    action_arg  = text,
                ))

        i += 1

    return {
        "question":     question,
        "steps":        steps,
        "final_answer": final_answer,
        "finished":     finished,
    }


# ── 5. LangChain AgentAction log ──────────────────────────────────────────────

def adapt_langchain_actions(actions_and_observations: list, question: str = "",
                             final_output: str = "") -> dict:
    """
    Normalize a LangChain agent trajectory (list of (AgentAction, observation) pairs).

    Input:
        actions_and_observations: list of
            ({"tool": "...", "tool_input": "...", "log": "Thought: ...\nAction: ..."}, "observation")
        question:     original user query
        final_output: the AgentFinish output string

    Handles both:
        - Dict-format actions (pre-v0.2): {"tool": ..., "tool_input": ..., "log": ...}
        - Object-format actions (v0.2+):  action.tool, action.tool_input, action.log
    """
    steps = []

    for item in actions_and_observations:
        if isinstance(item, (list, tuple)) and len(item) == 2:
            action, observation = item
        else:
            action      = item
            observation = ""

        # Support both dict and object forms
        if isinstance(action, dict):
            tool       = action.get("tool", "unknown")
            tool_input = action.get("tool_input", "")
            log        = action.get("log", "")
        else:
            tool       = getattr(action, "tool",       "unknown")
            tool_input = getattr(action, "tool_input", "")
            log        = getattr(action, "log",        "")

        # Extract thought from log (format: "Thought: ...\nAction: ...")
        thought = ""
        if log:
            m = re.search(r"Thought:\s*(.*?)(?:\nAction:|\Z)", log, re.DOTALL)
            if m:
                thought = m.group(1).strip()

        steps.append(_step(
            thought     = thought,
            action_type = _normalise_tool_name(tool),
            action_arg  = str(tool_input) if not isinstance(tool_input, str)
                          else tool_input,
            observation = str(observation),
        ))

    # Add Finish step
    if final_output:
        steps.append(_step(
            thought     = "",
            action_type = "Finish",
            action_arg  = final_output,
        ))

    return {
        "question":     question,
        "steps":        steps,
        "final_answer": final_output,
        "finished":     bool(final_output),
    }


# ── 6. Auto-detect dispatcher ─────────────────────────────────────────────────

def adapt_chain(raw: Any, question: str = "") -> dict:
    """
    Auto-detect the chain format and return a QPPG-normalised chain dict.

    Supported inputs
    ----------------
    - Native QPPG / ReAct dict  (has "steps" list with QPPG step keys)
    - OpenAI messages list       (list of {"role": ..., "content": ..., ...})
    - Structured tool-use dict   (has "turns" key)
    - LangChain trajectory list  (list of (AgentAction, observation) pairs)
    - OpenAI Assistants run      (dict with "run_steps" or "messages" key)

    Returns a dict with keys: question, steps, final_answer, finished
    """
    # ── Native QPPG dict ──────────────────────────────────────────────────────
    if isinstance(raw, dict) and "steps" in raw:
        steps = raw.get("steps", [])
        if steps and isinstance(steps[0], dict) and "action_type" in steps[0]:
            return adapt_react(raw)
        # Might be a structured output with "turns"
        if "turns" in raw:
            return adapt_structured_output(raw)

    # ── OpenAI messages list ──────────────────────────────────────────────────
    if isinstance(raw, list) and raw and isinstance(raw[0], dict):
        first = raw[0]
        # Function-call trace: assistant has tool_calls
        if any("tool_calls" in m for m in raw if isinstance(m, dict)):
            return adapt_openai_function_calls(raw, question=question)
        # Code-interpreter trace: content is a list of typed blocks
        if any(isinstance(m.get("content"), list) for m in raw if isinstance(m, dict)):
            return adapt_code_interpreter(raw, question=question)
        # LangChain trajectory: list of (action, obs) pairs or action dicts
        if isinstance(first, (list, tuple)):
            return adapt_langchain_actions(raw, question=question)
        # Plain messages list — treat as function-calling
        return adapt_openai_function_calls(raw, question=question)

    # ── Structured / tool-use dict ────────────────────────────────────────────
    if isinstance(raw, dict) and "turns" in raw:
        return adapt_structured_output(raw)

    # ── LangChain trajectory list ─────────────────────────────────────────────
    if isinstance(raw, list) and raw and isinstance(raw[0], (list, tuple)):
        return adapt_langchain_actions(raw, question=question)

    # Fallback: return empty chain
    return {
        "question":     question or str(raw)[:200],
        "steps":        [],
        "final_answer": "",
        "finished":     False,
    }


# ── Utility ───────────────────────────────────────────────────────────────────

def _normalise_tool_name(name: str) -> str:
    """
    Map diverse tool names to a canonical QPPG action_type.

    Examples:
        "web_search"      → "Search"
        "tavily_search"   → "Search"
        "python"          → "Code"
        "code_interpreter"→ "Code"
        "retriever"       → "Search"
        "calculator"      → "Tool"
    """
    name_lower = name.lower()
    if any(w in name_lower for w in ("search", "retriev", "lookup", "fetch", "browse")):
        return "Search"
    if any(w in name_lower for w in ("python", "code", "exec", "run", "interpret")):
        return "Code"
    if any(w in name_lower for w in ("finish", "final", "answer", "done", "complete")):
        return "Finish"
    return "Tool"
