"""
llm_agents.rl_agent
====================
LLM as a policy in a Gymnasium environment.

The agent uses chain-of-thought reasoning to select actions given the
observation, making it useful for environments where interpretability and
language-described action spaces matter.

Two modes
---------
1. LLMPolicy      – pure LLM policy (no learning, text-in text-out)
2. LLMAssistedQL  – tabular Q-learning where the LLM proposes an initial
                    Q-table bias from environment description (warm-start)

Usage – LLMPolicy
-----------------
    import gymnasium as gym
    from ml_ai_library.llm_agents.rl_agent import LLMPolicy

    env = gym.make("CartPole-v1")
    policy = LLMPolicy(
        model="gpt-4o-mini",
        env_description='''
        CartPole: balance a pole on a cart.
        Observation: [cart_pos, cart_vel, pole_angle, pole_ang_vel]
        Actions: 0=push left, 1=push right
        ''',
        action_meanings={0: "push left", 1: "push right"},
    )
    obs, _ = env.reset()
    for _ in range(200):
        action = policy.act(obs)
        obs, reward, done, *_ = env.step(action)
        if done: break

Usage – LLM Q-table warm-start
-------------------------------
    from ml_ai_library.llm_agents.rl_agent import LLMAssistedQL
    from ml_ai_library.norvig.rl import QLearning

    agent = LLMAssistedQL(
        model="gpt-4o-mini",
        env_description="FrozenLake 4x4 grid ...",
        n_actions=4,
        action_meanings={0:"L",1:"D",2:"R",3:"U"},
    )
    q_bias = agent.suggest_q_bias(state_examples=["(0,0)","(0,1)","(3,3)"])
"""

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

import numpy as np

from ml_ai_library.llm_agents.base_agent import BaseAgent


_POLICY_SYSTEM = """\
You are a reinforcement learning agent acting in an environment.

You will receive an observation and must choose the best action.
Think step-by-step, then output ONLY a JSON object:
  {"reasoning": "...", "action": <integer action index>}
"""

_QBIAS_SYSTEM = """\
You are an RL expert.  Given an environment description and a list of states,
estimate reasonable initial Q-values for each (state, action) pair.

Return JSON: {"q_values": {"<state>": [q0, q1, ...]}}
where qi is the estimated value for action i from that state.
Values should be in [-1, 1].
"""


class LLMPolicy:
    """
    Pure LLM policy – uses chain-of-thought to pick actions.

    Parameters
    ----------
    model             : litellm model string
    env_description   : text description of the environment
    action_meanings   : dict {action_int: meaning_str}
    observation_fn    : callable to convert obs → human-readable string.
                        Default: str(obs)
    """

    def __init__(
        self,
        model: str = "gpt-4o-mini",
        env_description: str = "",
        action_meanings: Optional[Dict[int, str]] = None,
        observation_fn=None,
        temperature: float = 0.0,
    ) -> None:
        self.action_meanings = action_meanings or {}
        self.obs_fn = observation_fn or str
        system = _POLICY_SYSTEM + f"\n\nENVIRONMENT:\n{env_description}"
        if action_meanings:
            am_str = "\n".join(f"  {k}: {v}" for k, v in action_meanings.items())
            system += f"\n\nACTION SPACE:\n{am_str}"
        self._agent = BaseAgent(model=model, system=system,
                                temperature=temperature)

    def act(self, observation: Any) -> int:
        """Select an action for the given observation."""
        obs_str = self.obs_fn(observation)
        self._agent.reset()
        reply   = self._agent.chat(f"Observation: {obs_str}\nChoose an action.")
        match   = re.search(r'"action"\s*:\s*(\d+)', reply)
        if match:
            return int(match.group(1))
        # Fallback: return 0
        return 0

    def evaluate(self, env, episodes: int = 5, max_steps: int = 200,
                 verbose: bool = False) -> List[float]:
        """Run the LLM policy for `episodes` episodes and return total rewards."""
        rewards = []
        for ep in range(episodes):
            obs, _ = env.reset()
            total  = 0.0
            for step in range(max_steps):
                action = self.act(obs)
                obs, reward, terminated, truncated, _ = env.step(action)
                total += reward
                if verbose:
                    print(f"  ep={ep} step={step} action={action} reward={reward}")
                if terminated or truncated:
                    break
            rewards.append(total)
        return rewards


class LLMAssistedQL:
    """
    Q-Learning warm-started by LLM Q-value suggestions.

    The LLM provides an initial Q-table bias for a set of example states;
    this can speed up convergence in sparse-reward environments.

    Parameters
    ----------
    model           : litellm model string
    env_description : text description of the environment
    n_actions       : number of discrete actions
    action_meanings : dict {action_int: meaning_str}
    """

    def __init__(
        self,
        model: str = "gpt-4o-mini",
        env_description: str = "",
        n_actions: int = 4,
        action_meanings: Optional[Dict[int, str]] = None,
        temperature: float = 0.0,
    ) -> None:
        self.n_actions = n_actions
        system = _QBIAS_SYSTEM + f"\n\nENVIRONMENT:\n{env_description}"
        if action_meanings:
            am_str = "\n".join(f"  {k}: {v}" for k, v in action_meanings.items())
            system += f"\n\nACTIONS:\n{am_str}"
        self._agent = BaseAgent(model=model, system=system,
                                temperature=temperature)

    def suggest_q_bias(self, state_examples: List[str]) -> Dict[str, List[float]]:
        """
        Ask the LLM for initial Q-value estimates for each example state.

        Returns
        -------
        dict {state_str: [q0, q1, ..., q_{n-1}]}
        """
        prompt = (
            f"States: {state_examples}\n"
            f"Provide initial Q-values (one list of {self.n_actions} floats per state)."
        )
        self._agent.reset()
        reply = self._agent.chat(prompt)
        match = re.search(r'"q_values"\s*:\s*(\{.*?\})', reply, re.DOTALL)
        if match:
            try:
                return json.loads(match.group(1))
            except json.JSONDecodeError:
                pass
        return {}
