"""
Robust parser — extract ASDP action_queue JSON from noisy LLM output.
Supports: raw JSON, ```json ... ```, { ... } fragments.
Allowed action types come from action_registry (protocol/actions.json).
"""
from __future__ import annotations

import json
import re
from typing import Any

from agentworld_gateway.action_registry import get_allowed_action_types

# Prefer ```json ... ``` or ``` ... ```
CODE_BLOCK = re.compile(r"```(?:json)?\s*([\s\S]*?)```", re.IGNORECASE)
# Match outermost { ... }, allow nesting
BRACE_OBJECT = re.compile(r"\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}")


def _parse_json_payload(text: str) -> dict[str, Any] | None:
    """Parse a JSON object (action_queue or internal_monologue) from text."""
    text = text.strip()
    try:
        return json.loads(text)
    except (json.JSONDecodeError, TypeError):
        pass
    for match in CODE_BLOCK.finditer(text):
        try:
            return json.loads(match.group(1).strip())
        except (json.JSONDecodeError, TypeError):
            continue
    for match in BRACE_OBJECT.finditer(text):
        try:
            return json.loads(match.group(0))
        except (json.JSONDecodeError, TypeError):
            continue
    return None


def extract_action_queue(raw: str) -> list[dict[str, Any]]:
    """
    Extract action_queue array from any text.
    Returns valid list[dict], or [{"type": "idle", "parameters": None}] on parse failure.
    """
    if not raw or not raw.strip():
        return [{"type": "idle", "parameters": None}]
    data = _parse_json_payload(raw.strip())
    return _normalize_actions(data) if data else [{"type": "idle", "parameters": None}]


def extract_action_queue_and_thought(raw: str) -> tuple[list[dict[str, Any]], str]:
    """
    Extract action_queue and internal_monologue from LLM reply.
    Returns (action_queue, internal_monologue).
    Uses short default if internal_monologue missing.
    """
    default_thought = "..."
    if not raw or not raw.strip():
        return [{"type": "idle", "parameters": None}], default_thought
    data = _parse_json_payload(raw.strip())
    if not data:
        return [{"type": "idle", "parameters": None}], default_thought
    if not isinstance(data, dict):
        if isinstance(data, list):
            actions = _normalize_actions({"action_queue": data})
        else:
            actions = [{"type": "idle", "parameters": None}]
        return actions, default_thought
    actions = _normalize_actions(data)
    thought = data.get("internal_monologue")
    if isinstance(thought, str) and thought.strip():
        thought = thought.strip()[:120]
    else:
        thought = default_thought
    return actions, thought


def _normalize_actions(data: dict[str, Any]) -> list[dict[str, Any]]:
    raw = data.get("action_queue") if isinstance(data, dict) else None
    if not isinstance(raw, list):
        return [{"type": "idle", "parameters": None}]
    out = []
    for item in raw:
        if not isinstance(item, dict):
            continue
        t = item.get("type")
        allowed = get_allowed_action_types()
        if t not in allowed:
            t = "idle"
        params = item.get("parameters")
        out.append({"type": t, "parameters": params})
    return out if out else [{"type": "idle", "parameters": None}]
