"""
OpenAI Assistants adapter for llm-guard-kit.

Parses an OpenAI Assistants API run's step history into QPPG chain format
and scores it with LabelFreeScorer or QppgMonitor.

Requirements:
    pip install "llm-guard-kit[openai]"
    # installs openai>=1.0

Usage:
    from openai import OpenAI
    from qppg_service.integrations.openai_adapter import score_assistants_run

    client = OpenAI()

    # After a run completes:
    result = score_assistants_run(client, thread_id, run_id)
    print(f"Risk: {result.risk_score:.3f}")

    # With a monitor for persistence:
    from qppg_service import QppgMonitor
    monitor = QppgMonitor(db_path="~/.qppg/chains.db")
    result  = score_assistants_run(client, thread_id, run_id, monitor=monitor)
"""

from __future__ import annotations

import json
from typing import Any, Optional

from ..label_free_scorer import LabelFreeScorer, LabelFreeResult
from ..monitor import QppgMonitor


def extract_question_from_thread(client: Any, thread_id: str) -> str:
    """
    Retrieve the first user message from the thread as the question string.

    Parameters
    ----------
    client    : openai.OpenAI instance
    thread_id : Assistants API thread ID

    Returns
    -------
    The content of the first user message, or empty string if not found.
    """
    try:
        messages = client.beta.threads.messages.list(
            thread_id=thread_id, order="asc", limit=1
        )
        for msg in messages.data:
            if msg.role == "user":
                content = msg.content
                if content and hasattr(content[0], "text"):
                    return content[0].text.value
    except Exception:
        pass
    return ""


def _extract_tool_step(tool_call: Any) -> dict:
    """Convert an OpenAI tool_call object into a QPPG step dict."""
    fn = getattr(tool_call, "function", None)
    if fn is not None:
        action_type = getattr(fn, "name", "tool_call")
        raw_args    = getattr(fn, "arguments", "")
        # Try to extract a clean string from JSON arguments
        try:
            args_dict = json.loads(raw_args)
            # Common argument key patterns: query, input, q, search_query
            for key in ("query", "input", "q", "search_query", "term"):
                if key in args_dict:
                    action_arg = str(args_dict[key])
                    break
            else:
                action_arg = raw_args[:200]
        except (json.JSONDecodeError, TypeError):
            action_arg = str(raw_args)[:200]
    else:
        # Retrieval or code_interpreter tool
        ttype = getattr(tool_call, "type", "retrieval")
        action_type = "Search" if ttype == "retrieval" else ttype
        action_arg  = ""

    # Normalise action_type
    lower = action_type.lower()
    if any(k in lower for k in ("search", "retriev", "lookup", "bing", "google", "tavily")):
        action_type = "Search"

    return {
        "thought":     "",
        "action_type": action_type,
        "action_arg":  action_arg,
        "observation": "",
    }


def score_assistants_run(
    client: Any,
    thread_id: str,
    run_id: str,
    scorer: Optional[LabelFreeScorer] = None,
    monitor: Optional[QppgMonitor] = None,
    question: str = "",
) -> LabelFreeResult:
    """
    Parse a completed OpenAI Assistants run and return a risk score.

    Parameters
    ----------
    client    : openai.OpenAI instance (must have beta.threads available)
    thread_id : Assistants API thread ID
    run_id    : Assistants API run ID
    scorer    : optional pre-built LabelFreeScorer (uses a new one if None)
    monitor   : optional QppgMonitor for persistence and auto-calibration
    question  : the user's question (auto-detected from thread if not given)

    Returns
    -------
    LabelFreeResult with risk_score and all component scores.
    """
    # Get question from thread if not provided
    if not question:
        question = extract_question_from_thread(client, thread_id)

    # Retrieve all run steps
    try:
        run_steps = client.beta.threads.runs.steps.list(
            thread_id=thread_id,
            run_id=run_id,
            order="asc",
        )
        steps_data = list(run_steps)
    except Exception as e:
        raise RuntimeError(f"Failed to retrieve run steps: {e}") from e

    chain_steps: list[dict] = []
    final_answer = ""
    finished = False

    for step in steps_data:
        step_type = getattr(step, "type", "")
        details   = getattr(step, "step_details", None)

        if step_type == "tool_calls" and details is not None:
            tool_calls = getattr(details, "tool_calls", [])
            for tc in tool_calls:
                s = _extract_tool_step(tc)
                chain_steps.append(s)

        elif step_type == "message_creation" and details is not None:
            # The assistant's final message
            msg_id = getattr(getattr(details, "message_creation", None), "message_id", None)
            if msg_id:
                try:
                    msg = client.beta.threads.messages.retrieve(
                        thread_id=thread_id, message_id=msg_id
                    )
                    content = getattr(msg, "content", [])
                    if content and hasattr(content[0], "text"):
                        final_answer = content[0].text.value
                        finished = True
                except Exception:
                    pass
            # Add a Finish step
            chain_steps.append({
                "thought":     "",
                "action_type": "Finish",
                "action_arg":  final_answer,
                "observation": "",
            })

    # Fill observations: the observation for step[i] is usually the result of
    # the tool call, which the Assistants API exposes as the step's output.
    # If retrieval steps have output, attach it.
    for i, step in enumerate(steps_data):
        step_type = getattr(step, "type", "")
        if step_type == "tool_calls":
            details = getattr(step, "step_details", None)
            tool_calls = getattr(details, "tool_calls", []) if details else []
            for j, tc in enumerate(tool_calls):
                # Retrieval results are not directly exposed in run steps;
                # use the type hint if available
                ttype = getattr(tc, "type", "")
                if ttype == "retrieval":
                    # Observation not available in Assistants API — mark as retrieved
                    idx = sum(
                        len(getattr(getattr(s, "step_details", None), "tool_calls", []))
                        for s in steps_data[:i]
                        if getattr(s, "type", "") == "tool_calls"
                    ) + j
                    if idx < len(chain_steps):
                        chain_steps[idx]["observation"] = "[retrieved]"

    scorer = scorer or (monitor.scorer if monitor else LabelFreeScorer())
    chain = {
        "question":     question,
        "steps":        chain_steps,
        "final_answer": final_answer,
        "finished":     finished,
    }

    if monitor is not None:
        monitor.track(question, chain_steps, final_answer, finished=finished)
        # Return the score from the scorer directly
        return scorer.score(question, chain_steps, final_answer, finished)

    return scorer.score(question, chain_steps, final_answer, finished)
