"""
llm_agents.planning_agent
==========================
LLM-powered planner that emits STRIPS-style or HTN task networks.

The agent is given a domain description and a goal; it uses chain-of-thought
to decompose the goal into an ordered plan, then optionally validates the plan
against a STRIPSProblem or HTNProblem.

Usage
-----
    from ml_ai_library.llm_agents.planning_agent import PlanningAgent

    agent = PlanningAgent(model="gpt-4o-mini")
    plan  = agent.plan(
        domain='''
        Actions available:
          - pick_up(block): requires clear(block) and on_table(block)
          - put_down(block): block is in hand
          - stack(block, on): requires holding(block) and clear(on)
          - unstack(block, from): requires on(block, from) and clear(block)
        ''',
        initial_state="on_table(A), on_table(B), clear(A), clear(B)",
        goal="on(A, B)",
    )
    print(plan)   # ["pick_up(A)", "stack(A, B)"]

Validation
----------
    from ml_ai_library.norvig.planning import STRIPSProblem, Action
    strips_prob = STRIPSProblem(...)
    valid_plan  = agent.plan_and_validate(domain, init, goal, strips_prob)
"""

from __future__ import annotations
import json
import re
from typing import List, Optional

from ml_ai_library.llm_agents.base_agent import BaseAgent

_PLAN_SYSTEM = """\
You are an AI planning assistant.

Given a domain description, an initial state, and a goal, produce a plan:
a JSON list of action strings that transforms the initial state into the goal.

Return ONLY valid JSON of the form:
  {"plan": ["action1", "action2", ...], "rationale": "brief explanation"}

Think step by step before outputting.
"""


class PlanningAgent:
    """
    LLM planner that generates STRIPS-style action sequences.

    Parameters
    ----------
    model : litellm model string
    """

    def __init__(self, model: str = "gpt-4o-mini", temperature: float = 0.0):
        self._agent = BaseAgent(
            model=model,
            system=_PLAN_SYSTEM,
            temperature=temperature,
        )

    def plan(
        self,
        domain: str,
        initial_state: str,
        goal: str,
    ) -> List[str]:
        """
        Generate a plan as a list of action strings.

        Parameters
        ----------
        domain        : natural-language domain description (actions + preconditions)
        initial_state : description of the initial world state
        goal          : description of the goal state

        Returns
        -------
        List of action strings
        """
        prompt = (
            f"DOMAIN:\n{domain}\n\n"
            f"INITIAL STATE:\n{initial_state}\n\n"
            f"GOAL:\n{goal}\n\n"
            "Produce a plan as JSON."
        )
        self._agent.reset()
        reply = self._agent.chat(prompt)

        # Extract JSON block
        match = re.search(r"\{.*\}", reply, re.DOTALL)
        if match:
            data = json.loads(match.group())
            return data.get("plan", [])
        return []

    def plan_and_validate(
        self,
        domain: str,
        initial_state: str,
        goal: str,
        strips_problem,
        max_retries: int = 3,
    ) -> Optional[List[str]]:
        """
        Generate a plan and validate it against a STRIPSProblem.
        Retries up to max_retries times if validation fails.
        """
        from ml_ai_library.norvig.planning import STRIPSProblem

        for attempt in range(max_retries):
            plan_strs = self.plan(domain, initial_state, goal)
            # Try to match plan strings to STRIPSProblem actions by name
            action_map = {a.name: a for a in strips_problem.actions}
            actions = [action_map[s] for s in plan_strs if s in action_map]
            # Simulate the plan
            state = strips_problem.initial
            valid = True
            for action in actions:
                if not action.applicable(state):
                    valid = False
                    break
                state = action.apply(state)
            if valid and strips_problem.goal <= state:
                return plan_strs
            # Add failure feedback and retry
            self._agent.chat(
                f"That plan was invalid at step {len(actions)}. "
                "Please try again with a correct plan."
            )
        return None
