"""
openai_baseline.py — OpenAI LLM baseline for the Energizer environment.

This script uses the OpenAI chat completions API to run a language model as a
battery dispatch agent. The LLM receives a structured plain-English description
of the current environment state at each 5-minute step and must reply with a
single JSON object containing the power action.

API Key Resolution Order:
  1. OPENAI_API_KEY environment variable (recommended for CI/Docker).
  2. Interactive terminal prompt (fallback for local use without a key set).

The existing DDPG baselines in agent/baselines.py are NOT touched or removed.
This file is an independent, additive compliance script.

Usage:
  python openai_baseline.py                        # auto-detects key
  OPENAI_API_KEY=sk-... python openai_baseline.py  # explicit key
"""

import os
import json
import getpass
import numpy as np
import sys

# Add root directory to sys.path for absolute imports
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from openenv.grid_twin import GridTwinEnv
from openenv.models import Action
from openenv.tasks import normalize_score

# ──────────────────────────────────────────────
# 1. RESOLVE API KEY
# ──────────────────────────────────────────────

def get_api_key() -> str:
    """
    Returns the OpenAI API key.
    Checks OPENAI_API_KEY env var first; falls back to interactive prompt.
    Uses getpass so the key is never echoed to the terminal.
    """
    key = os.environ.get("OPENAI_API_KEY", "").strip()
    if key:
        print("[openai_baseline] API key loaded from environment variable.")
        return key

    print("\n[openai_baseline] OPENAI_API_KEY not found in environment variables.")
    print("Please enter your OpenAI API key below (input is hidden):")
    key = getpass.getpass("OpenAI API Key: ").strip()
    if not key:
        raise ValueError(
            "No API key provided. Set OPENAI_API_KEY or enter it interactively."
        )
    return key


# ──────────────────────────────────────────────
# 2. SYSTEM PROMPT
# ──────────────────────────────────────────────

SYSTEM_PROMPT = """You are an expert battery energy storage controller.
Your goal is to maximize profit by controlling a grid-connected battery across three services:

1. Energy Arbitrage (EA): Buy (charge) when electricity prices (LMP) are low, sell (discharge) when prices are high.
2. Frequency Regulation (FR): Follow the RegD signal. Positive RegD means you should charge, negative means discharge.
3. Peak Shaving (PS): Discharge when demand is high (near or above peak threshold) to reduce grid stress.

Battery constraints:
- State of Charge (SOC) must stay between 0.0 and 1.0.
- Power action must be in [-1.0, 1.0] where +1.0 is full charge and -1.0 is full discharge.
- Avoid keeping SOC near the extremes (0 or 1) as it limits your flexibility.

You will receive the current state of the environment. Respond with ONLY a valid JSON object:
{"power": <float between -1.0 and 1.0>}

No explanation. No markdown. Just the JSON object."""


# ──────────────────────────────────────────────
# 3. STATE → PROMPT
# ──────────────────────────────────────────────

def build_user_prompt(state: dict) -> str:
    """
    Converts the env.state() dict into a structured plain-English prompt.
    All values are rounded for token efficiency. The LMP average is inferred
    from the current LMP to keep the prompt self-contained.
    """
    steps_remaining = state["T"] - state["t"]
    steps_to_peak = (state["peak_time"] - state["t"]) % state["T"]

    return f"""Current battery environment state:
- State of Charge (SOC): {state['soc']:.3f}  (0=empty, 1=full)
- Current electricity price (LMP): {state['lmp_current']:.3f}
- Frequency regulation signal (RegD): {state['regd_current']:.3f}  (follow this signal)
- Demand (normalised by peak): {state['demand_current']:.3f}  (>1.0 means above peak threshold)
- Peak demand threshold: {state['peak_threshold']:.3f}
- Steps to peak demand: {steps_to_peak} (out of {state['T']})
- Steps remaining in episode: {steps_remaining}
- Cumulative EA profit so far: {state['total_EA']:.2f}
- Cumulative FR profit so far: {state['total_FR']:.2f}
- Cumulative PS profit so far: {state['total_PS']:.2f}

What power action should the battery take right now?"""


# ──────────────────────────────────────────────
# 4. LLM AGENT
# ──────────────────────────────────────────────

def llm_policy_factory(client, model: str = "gpt-4o-mini"):
    """
    Returns a policy function that calls the OpenAI API for each step.
    The factory captures `client` and `model` in its closure; the returned
    function matches the standard signature: policy(obs_array) -> float.

    Note: obs_array is the numpy observation but we call env.state() externally
    for the richer prompt. This function receives obs_array for compatibility
    with the grade() interface but the actual env state is fetched inside
    the episode loop (see run_llm_episode).
    """
    def policy(obs_array):
        # This stub exists for interface compatibility.
        # Actual LLM calls happen in run_llm_episode() where we have env access.
        return 0.0
    return policy


def query_llm(client, model: str, state: dict) -> float:
    """
    Sends the current state to the LLM and parses the JSON action.
    Returns power as a float clamped to [-1.0, 1.0].
    Fails gracefully: returns 0.0 (do nothing) if parsing fails.
    """
    user_prompt = build_user_prompt(state)

    try:
        response = client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt},
            ],
            temperature=0.0,   # deterministic for reproducibility
            max_tokens=50,     # we only need {"power": X}
        )
        text = response.choices[0].message.content.strip()
        action_dict = json.loads(text)
        power = float(action_dict["power"])
        return float(np.clip(power, -1.0, 1.0))

    except (json.JSONDecodeError, KeyError, ValueError) as e:
        print(f"  [warn] LLM response parse error: {e}. Defaulting to 0.0.")
        return 0.0


# ──────────────────────────────────────────────
# 5. EPISODE RUNNER
# ──────────────────────────────────────────────

def run_llm_episode(client, model: str, difficulty: str, seed: int) -> float:
    """
    Runs a single episode of the environment using the LLM as the agent.
    Uses env.state() for the rich prompt and env.step(Action(power=x)) for control.
    Returns the total undiscounted reward for the episode.
    """
    from openenv.tasks import Task

    task = Task(difficulty)
    task.env = GridTwinEnv(seed=seed)
    task.env.reset(episode=999)

    # Apply difficulty modifications (mirrors tasks.py exactly)
    if difficulty == "easy":
        task.env.regd_series *= 0.5
        task.env.demand_series *= 0.7
    elif difficulty == "hard":
        rng = task.env.rng
        task.env.lmp_series += rng.normal(0, 0.5, task.env.T)
        task.env.regd_series = np.clip(
            task.env.regd_series + rng.normal(0, 0.4, task.env.T),
            -1, 1
        )

    total_reward = 0.0
    done = False

    while not done:
        idx = min(task.env.t, task.env.T - 1)
        state = {
            "soc": task.env.soc,
            "lmp_current": float(task.env.lmp_series[idx]),
            "regd_current": float(task.env.regd_series[idx]),
            "demand_current": float(task.env.demand_series[idx]) / (task.env.peak_threshold + 1e-6),
            "peak_threshold": float(task.env.peak_threshold),
            "peak_time": int(np.argmax(task.env.demand_series)),
            "T": task.env.T,
            "t": task.env.t,
            "total_EA": float(task.env.total_EA),
            "total_FR": float(task.env.total_FR),
            "total_PS": float(task.env.total_PS)
        }
        power = query_llm(client, model, state)
        _, reward, done, _ = task.env.step(Action(power=power))
        total_reward += reward

    return total_reward


# ──────────────────────────────────────────────
# 6. GRADING
# ──────────────────────────────────────────────

def grade_llm(client, model: str = "gpt-4o-mini", num_episodes: int = 3):
    """
    Runs the LLM agent across all 3 difficulty tiers and prints normalised scores.
    num_episodes is intentionally kept small (default 3) to limit API costs.
    Each episode call to the LLM makes T=288 API requests, so 3 episodes per
    difficulty = 864 requests per tier, 2592 total.

    Returns a dict: {"easy": float, "medium": float, "hard": float, "final": float}
    """
    results = {}
    print(f"\n=== OpenAI LLM Baseline Grading (model={model}, episodes={num_episodes}) ===\n")

    for difficulty in ["easy", "medium", "hard"]:
        rewards = []
        for i in range(num_episodes):
            print(f"  [{difficulty.upper()}] Episode {i + 1}/{num_episodes}...", end=" ", flush=True)
            r = run_llm_episode(client, model, difficulty, seed=i)
            rewards.append(r)
            print(f"reward={r:.2f}")

        avg_reward = float(np.mean(rewards))
        score = normalize_score(avg_reward, difficulty)
        results[difficulty] = score
        print(f"  ↳ Avg reward={avg_reward:.2f}  Normalised score={score:.4f}\n")

    results["final"] = float(np.mean([results["easy"], results["medium"], results["hard"]]))
    print(f"Final LLM Baseline Score: {results['final']:.4f}")
    return results


# ──────────────────────────────────────────────
# 7. MAIN
# ──────────────────────────────────────────────

if __name__ == "__main__":
    # Resolve API key
    api_key = get_api_key()

    # Import openai here so the rest of the module is importable even without it
    try:
        from openai import OpenAI
    except ImportError:
        print("\n[error] The 'openai' package is not installed.")
        print("Run: pip install openai")
        raise

    client = OpenAI(api_key=api_key)

    # Use gpt-4o-mini by default (cheap, fast, capable enough for structured JSON)
    MODEL = os.environ.get("OPENAI_MODEL", "gpt-4o-mini")

    # num_episodes=1 by default for a quick test; increase for a stable estimate
    NUM_EPISODES = int(os.environ.get("OPENAI_BASELINE_EPISODES", "1"))

    print(f"[openai_baseline] Using model: {MODEL}")
    print(f"[openai_baseline] Episodes per difficulty: {NUM_EPISODES}")

    scores = grade_llm(client, model=MODEL, num_episodes=NUM_EPISODES)
    print("\nFull results:", scores)