"""Node that calls OpenAI LLM to generate NPC responses."""

import json
import re

from langchain_core.messages import AIMessage
from langchain_openai import ChatOpenAI

from src.schemas.state import AgentOutput, WalkState


def _get_llm():
    return ChatOpenAI(
        model="gpt-4o-mini",
        temperature=0.7,
        model_kwargs={"response_format": {"type": "json_object"}},
    )


def _extract_json(text: str) -> dict:
    """Try multiple strategies to extract JSON from LLM response."""

    text = text.strip()

    # 1) Strip code block wrappers
    if text.startswith("```"):
        text = re.sub(r"^```(?:json)?\s*", "", text)
        text = re.sub(r"\s*```$", "", text)
        text = text.strip()

    # 2) Try direct parse
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass

    # 3) Find first { ... } block
    match = re.search(r"\{.*\}", text, re.DOTALL)
    if match:
        try:
            return json.loads(match.group())
        except json.JSONDecodeError:
            pass

    return {}


def _clamp(value: float, lo: float, hi: float) -> float:
    return max(lo, min(hi, value))


def call_llm(state: WalkState) -> dict:
    """Call LLM and parse the AgentOutput with extended fields."""

    messages = state.get("messages", [])

    try:
        result = _get_llm().invoke(messages)
        parsed = _extract_json(result.content)

        if not parsed.get("response"):
            parsed = {
                "response": result.content.strip(),
                "should_end_turn": False,
                "next_step": "continue",
                "emotion_assessment": {"valence": 0.0, "arousal": 0.0, "label": "neutral"},
                "observation": "",
                "empathy_quality": "neutral",
                "trust_shift": 0.0,
                "npc_internal_mood": "guarded",
            }

        # Validate and clamp numeric fields
        emotion = parsed.get("emotion_assessment", {})
        trust_shift = _clamp(float(parsed.get("trust_shift", 0.0)), -0.3, 0.3)

        empathy_quality = parsed.get("empathy_quality", "neutral")
        if empathy_quality not in ("hit", "miss", "neutral"):
            empathy_quality = "neutral"

        npc_mood = parsed.get("npc_internal_mood", "guarded")
        if npc_mood not in ("guarded", "cautious", "warming", "open", "withdrawn"):
            npc_mood = "guarded"

        agent_output: AgentOutput = {
            "response": parsed.get("response", ""),
            "should_end_turn": bool(parsed.get("should_end_turn", False)),
            "next_step": parsed.get("next_step", "continue"),
            "emotion_assessment": {
                "valence": _clamp(float(emotion.get("valence", 0.0)), -1.0, 1.0),
                "arousal": _clamp(float(emotion.get("arousal", 0.0)), 0.0, 1.0),
                "label": emotion.get("label", "neutral"),
            },
            "observation": parsed.get("observation", ""),
            "empathy_quality": empathy_quality,
            "trust_shift": trust_shift,
            "npc_internal_mood": npc_mood,
        }

        # Add AI message to history
        new_messages = list(messages) + [AIMessage(content=agent_output["response"])]

        return {
            "agent_output": agent_output,
            "messages": new_messages,
        }

    except Exception as e:
        return {"error": f"LLM call failed: {e}"}