"""Natural-language policy (P1.3): parses LLM JSON replies into Action sequences.

P5.2 adds a bounded repair loop: on a parse failure the policy re-asks the
model with an accumulated error trail up to ``max_repairs`` times, then
falls back to the scripted policy; ``NLPolicyStats`` exposes the call and
repair counters for the sweep and the phase-3 report.
"""

from __future__ import annotations

import json
from collections.abc import Callable
from dataclasses import dataclass, field, fields
from typing import TYPE_CHECKING, Literal, TypeGuard, get_args

from pf1e_simulator.candidates import candidates
from pf1e_simulator.combat import Action, ReadyTrigger, TrickCondition

if TYPE_CHECKING:
    from pf1e_simulator.client import LLMClient
    from pf1e_simulator.combat import Observation
    from pf1e_simulator.prompt_builder import PromptBuilder


# ---- Type aliases for the 27 supported kinds (filtered from Action.kind) ----
KindValue = Literal[
    "attack", "move", "wait", "immediate", "full_attack", "charge", "withdraw",
    "5foot_step", "cast_spell", "maneuver", "channel_energy", "hex", "deed",
    "spell_combat", "inspire_courage", "retrieve_weapon", "repair_weapon",
    "activate_power_attack", "delay", "ready", "total_defense",
    "fight_defensively", "stand_up", "drop_prone", "aid_another", "run", "hide",
]
_KIND_VALUES: frozenset[str] = frozenset(get_args(KindValue))
_FIELD_NAMES: frozenset[str] = frozenset(f.name for f in fields(Action))
_PAIR_LEN = 2
_CHANNEL_MODES: frozenset[str] = frozenset(get_args(Literal["heal", "harm"]))
_AID_MODES: frozenset[str] = frozenset(get_args(Literal["attack", "ac"]))
_TRICK_CONDITIONS: frozenset[str] = frozenset(get_args(TrickCondition))
_READY_TRIGGERS: frozenset[str] = frozenset(get_args(ReadyTrigger))


class NLPolicyError(ValueError):
    """The LLM reply is not a valid JSON action sequence."""


# ---- Type guards -------------------------------------------------------------

def _is_kind(value: object) -> TypeGuard[KindValue]:
    return value in _KIND_VALUES


def _is_channel_mode(value: object) -> TypeGuard[Literal["heal", "harm"]]:
    return value in _CHANNEL_MODES


def _is_aid_mode(value: object) -> TypeGuard[Literal["attack", "ac"]]:
    return value in _AID_MODES


def _is_trick_condition(value: object) -> TypeGuard[TrickCondition]:
    return value in _TRICK_CONDITIONS


def _is_ready_trigger(value: object) -> TypeGuard[ReadyTrigger]:
    return value in _READY_TRIGGERS


def _is_object_list(value: object) -> TypeGuard[list[object]]:
    return isinstance(value, list)


def _is_object_seq(value: object) -> TypeGuard[list[object] | tuple[object, ...]]:
    return isinstance(value, (list, tuple))


def _is_object_dict(value: object) -> TypeGuard[dict[str, object]]:
    return isinstance(value, dict)


# ---- Coercion helpers --------------------------------------------------------

def _opt_str(value: object, name: str) -> str | None:
    if value is None:
        return None
    if isinstance(value, str):
        return value
    msg = f"{name!r} must be a string"
    raise NLPolicyError(msg)


def _str_tuple(value: object, name: str) -> tuple[str, ...]:
    if value is None:
        return ()
    if not _is_object_seq(value):
        msg = f"{name!r} must be a list of strings"
        raise NLPolicyError(msg)
    out: list[str] = []
    for item in value:
        if not isinstance(item, str):
            msg = f"{name!r} must be a list of strings"
            raise NLPolicyError(msg)
        out.append(item)
    return tuple(out)


def _int_pair(value: object, name: str) -> tuple[int, int] | None:
    if value is None:
        return None
    if not _is_object_seq(value):
        msg = f"{name!r} must be a [row, col] pair"
        raise NLPolicyError(msg)
    if len(value) != _PAIR_LEN:
        msg = f"{name!r} must be a [row, col] pair"
        raise NLPolicyError(msg)
    first, second = value[0], value[1]
    if (
        isinstance(first, bool)
        or isinstance(second, bool)
        or not isinstance(first, int)
        or not isinstance(second, int)
    ):
        msg = f"{name!r} must be a [row, col] pair"
        raise NLPolicyError(msg)
    return (first, second)


def _opt_bool(value: object, name: str) -> bool:
    if value is None:
        return False
    if isinstance(value, bool):
        return value
    msg = f"{name!r} must be a boolean"
    raise NLPolicyError(msg)


def _opt_channel_mode(value: object) -> Literal["heal", "harm"] | None:
    if value is None:
        return None
    if _is_channel_mode(value):
        return value
    msg = "'channel_mode' must be one of: heal, harm"
    raise NLPolicyError(msg)


def _opt_aid_mode(value: object) -> Literal["attack", "ac"] | None:
    if value is None:
        return None
    if _is_aid_mode(value):
        return value
    msg = "'aid_mode' must be one of: attack, ac"
    raise NLPolicyError(msg)


def _opt_trick_condition(value: object) -> TrickCondition | None:
    if value is None:
        return None
    if _is_trick_condition(value):
        return value
    msg = "'trick_condition' must be one of: " + ", ".join(get_args(TrickCondition))
    raise NLPolicyError(msg)


def _opt_ready_trigger(value: object) -> ReadyTrigger | None:
    if value is None:
        return None
    if _is_ready_trigger(value):
        return value
    msg = "'ready_trigger' must be one of: " + ", ".join(get_args(ReadyTrigger))
    raise NLPolicyError(msg)


# ---- Action construction -----------------------------------------------------

def _action_from_item(item: object) -> Action:
    if not _is_object_dict(item):
        msg = "each action must be a JSON object"
        raise NLPolicyError(msg)
    fields_dict = item
    unknown = sorted(set(fields_dict) - _FIELD_NAMES)
    if unknown:
        msg = f"unknown action field(s): {', '.join(unknown)}"
        raise NLPolicyError(msg)
    kind = fields_dict.get("kind")
    if not _is_kind(kind):
        msg = "missing or unknown 'kind'"
        raise NLPolicyError(msg)
    return Action(
        kind=kind,
        target_id=_opt_str(fields_dict.get("target_id"), "target_id"),
        spell_name=_opt_str(fields_dict.get("spell_name"), "spell_name"),
        maneuver=_opt_str(fields_dict.get("maneuver"), "maneuver"),
        hex_name=_opt_str(fields_dict.get("hex_name"), "hex_name"),
        deed_name=_opt_str(fields_dict.get("deed_name"), "deed_name"),
        class_ability=_opt_str(fields_dict.get("class_ability"), "class_ability"),
        channel_mode=_opt_channel_mode(fields_dict.get("channel_mode")),
        channel_exclude=_str_tuple(fields_dict.get("channel_exclude"), "channel_exclude"),
        direction=_int_pair(fields_dict.get("direction"), "direction"),
        trick_condition=_opt_trick_condition(fields_dict.get("trick_condition")),
        weapon_id=_opt_str(fields_dict.get("weapon_id"), "weapon_id"),
        ready_trigger=_opt_ready_trigger(fields_dict.get("ready_trigger")),
        destination=_int_pair(fields_dict.get("destination"), "destination"),
        target_pos=_int_pair(fields_dict.get("target_pos"), "target_pos"),
        aid_mode=_opt_aid_mode(fields_dict.get("aid_mode")),
        metamagic=_str_tuple(fields_dict.get("metamagic"), "metamagic"),
        cast_defensively=_opt_bool(fields_dict.get("cast_defensively"), "cast_defensively"),
        tumble=_opt_bool(fields_dict.get("tumble"), "tumble"),
    )


# ---- Public parse API --------------------------------------------------------

def parse_actions(reply: str) -> tuple[Action, ...]:
    """Parse an LLM reply into a tuple of Action objects."""
    try:
        parsed = json.loads(reply)
    except json.JSONDecodeError as exc:
        msg = f"reply is not valid JSON: {exc}"
        raise NLPolicyError(msg) from exc
    if not _is_object_list(parsed):
        msg = "reply must be a JSON array of action objects"
        raise NLPolicyError(msg)
    items = parsed
    return tuple(_action_from_item(item) for item in items)


# ---- Fallback policy ---------------------------------------------------------

Fallback = Callable[[tuple[Action, ...]], tuple[Action, ...]]


def fallback_default(menu: tuple[Action, ...]) -> tuple[Action, ...]:
    """Deterministic fallback: first decisive action, then a defensive hold."""
    for kind in ("full_attack", "attack", "charge", "move", "5foot_step"):
        for action in menu:
            if action.kind == kind:
                return (action,)
    if menu:
        return (menu[0],)
    return (Action(kind="wait"),)


# ---- Repair protocol ---------------------------------------------------------

_REPAIR_INSTRUCTION = "Return a valid JSON action sequence."
_REPAIR_HEADER_TEMPLATE = "-- previous reply invalid (attempt {}) --\n{}"


@dataclass
class NLPolicyStats:
    """Mutable counters for the LLM decision loop (repair observability).

    The phase-3 report and the sweep pipeline read these counters to
    quantify how often the LLM layer repairs invalid replies or falls back
    to the scripted policy.
    """

    calls: int = 0
    repairs_attempted: int = 0
    repairs_succeeded: int = 0
    fallbacks: int = 0


# ---- NLPolicy ----------------------------------------------------------------

@dataclass(frozen=True)
class NLPolicy:
    """Turn LLM completions into Action sequences, with a scripted fallback."""

    client: LLMClient
    builder: PromptBuilder
    prompt_name: str = "nl_default"
    fallback: Fallback = fallback_default
    max_repairs: int = 1
    stats: NLPolicyStats = field(default_factory=NLPolicyStats)

    def __post_init__(self) -> None:
        if self.max_repairs < 0:
            msg = f"max_repairs must be >= 0, got {self.max_repairs}"
            raise ValueError(msg)
        expected = self.builder.prompt_version(self.prompt_name)
        if self.client.prompt_version != expected:
            msg = (
                f"client prompt_version {self.client.prompt_version!r} does not "
                f"match pack {expected!r} for {self.prompt_name!r}"
            )
            raise ValueError(msg)

    def __call__(self, observation: Observation) -> tuple[Action, ...]:
        menu = candidates(observation)
        prompt = self.builder.build(observation, prompt_name=self.prompt_name)
        reply = self.client.complete(prompt)
        self.stats.calls += 1
        errors: list[NLPolicyError] = []
        while True:
            try:
                actions = parse_actions(reply)
                break
            except NLPolicyError as exc:
                if len(errors) >= self.max_repairs:
                    self.stats.fallbacks += 1
                    return self.fallback(menu)
                errors.append(exc)
                self.stats.repairs_attempted += 1
                trail = "\n\n".join(
                    _REPAIR_HEADER_TEMPLATE.format(index, error)
                    for index, error in enumerate(errors, start=1)
                )
                retry_prompt = f"{prompt}\n\n{trail}\n{_REPAIR_INSTRUCTION}"
                reply = self.client.complete(retry_prompt)
                self.stats.calls += 1
        if errors:
            self.stats.repairs_succeeded += 1
        return actions