"""
inference.py — Baseline inference script
Uses OpenAI client with API_BASE_URL / MODEL_NAME / HF_TOKEN env vars.
Must run in <20 min on 2 vCPU / 8GB.
"""
import os, sys, json
sys.path.insert(0, "server")   # so grader/environment are importable

import requests
from openai import OpenAI
from server.environment import SensitivityAction

API_BASE_URL = os.environ.get("API_BASE_URL", "https://api.openai.com/v1")
MODEL_NAME   = os.environ.get("MODEL_NAME",   "gpt-4o-mini")
HF_TOKEN     = os.environ.get("HF_TOKEN",     "")
ENV_URL      = os.environ.get("ENV_URL",      "http://localhost:7860")

client = OpenAI(api_key=HF_TOKEN, base_url=API_BASE_URL)


def env_reset():
    r = requests.post(f"{ENV_URL}/reset", timeout=30)
    r.raise_for_status()
    return r.json()

def env_step(action: dict):
    r = requests.post(f"{ENV_URL}/step", json={"action": action}, timeout=30)
    r.raise_for_status()
    return r.json()

def build_prompt(obs: dict) -> str:
    variants  = "\n".join(f"  [{i}] \"{v}\"" for i, v in enumerate(obs["variants"]))
    responses = "\n".join(f"  Response [{i}]: \"{r}\"" for i, r in enumerate(obs["ai_responses"]))
    return f"""You are an expert AI evaluator specializing in prompt sensitivity analysis.

Base Prompt: "{obs['base_prompt']}"

Prompt Variants:
{variants}

AI Responses to each variant:
{responses}

Task: {obs['instruction']}

Respond ONLY with a valid JSON object — no markdown, no extra text:
{{
  "verdict": "<sensitive|stable|partial>",
  "confidence": <float 0.0-1.0>,
  "explanation": "<your reasoning, minimum 20 characters>",
  "sensitive_variant_index": <integer index or null>
}}"""

def call_agent(obs: dict) -> dict:
    raw = client.chat.completions.create(
        model=MODEL_NAME,
        messages=[{"role": "user", "content": build_prompt(obs)}],
        temperature=0.0,
    ).choices[0].message.content.strip()
    if raw.startswith("```"):
        raw = raw.split("```")[1]
        if raw.startswith("json"):
            raw = raw[4:]
    return json.loads(raw.strip())


def run(num_episodes: int = 5):
    print(f"\n{'='*58}")
    print(f"  NeuroHack — Task 1: Prompt Sensitivity Baseline")
    print(f"  Model  : {MODEL_NAME}")
    print(f"  Env    : {ENV_URL}")
    print(f"{'='*58}\n")

    scores = []
    for ep in range(1, num_episodes + 1):
        print(f"── Episode {ep}/{num_episodes} ──────────────────────────")
        data = env_reset()
        obs  = data["observation"]
        print(f"   Base prompt : {obs['base_prompt']}")

        try:
            action_dict = call_agent(obs)
        except Exception as e:
            print(f"   [AGENT ERROR] {e} — using fallback")
            action_dict = {
                "verdict": "stable",
                "confidence": 0.5,
                "explanation": "Agent error — fallback to stable verdict.",
                "sensitive_variant_index": None,
            }

        print(f"   Verdict     : {action_dict.get('verdict')}")
        print(f"   Confidence  : {action_dict.get('confidence')}")
        print(f"   Explanation : {str(action_dict.get('explanation',''))[:80]}")

        result     = env_step(action_dict)
        reward     = result["reward"]
        info       = result["info"]
        scores.append(reward)

        print(f"   Reward      : {reward}")
        print(f"   GT verdict  : {info.get('ground_truth')}")
        print(f"   Breakdown   : {info.get('breakdown')}\n")

    avg = round(sum(scores) / len(scores), 4)
    print(f"{'='*58}")
    print(f"  Episodes : {num_episodes}")
    print(f"  Scores   : {scores}")
    print(f"  Average  : {avg}")
    print(f"{'='*58}\n")
    return scores, avg


if __name__ == "__main__":
    run(num_episodes=5)