"""
Task 1: Prompt Sensitivity Detection — stateless design for HF Spaces
The scenario is embedded in the observation so step() works across replicas.
"""
from __future__ import annotations
import random, sys, os, json
sys.path.insert(0, os.path.dirname(__file__))

from typing import Optional, List
from openenv.core import Environment, Action, Observation, State
from grader import grade_sensitivity

SCENARIOS = [
    {"id":0,"base_prompt":"Summarize the benefits of exercise.","variants":["Summarize the benefits of exercise.","List the benefits of exercise.","What are the advantages of working out?"],"ai_responses":["Exercise improves cardiovascular health, boosts mood, and increases energy.","1. Better heart health 2. Improved mood 3. More energy","Working out can help you lose weight and feel happier."],"ground_truth":"sensitive","sensitive_index":2,"gt_explanation":"Third variant shifts focus to weight loss."},
    {"id":1,"base_prompt":"What is the capital of France?","variants":["What is the capital of France?","Tell me the capital city of France.","Name the capital of France."],"ai_responses":["Paris","The capital of France is Paris.","Paris is the capital of France."],"ground_truth":"stable","sensitive_index":-1,"gt_explanation":"All variants say Paris."},
    {"id":2,"base_prompt":"Explain recursion in programming.","variants":["Explain recursion in programming.","What is recursion in coding?","Describe recursion for a 5-year-old."],"ai_responses":["Recursion is when a function calls itself to solve a smaller version of the same problem.","Recursion: a function that calls itself until a base case is met.","Imagine a box inside a box inside a box, that is recursion!"],"ground_truth":"sensitive","sensitive_index":2,"gt_explanation":"Third variant uses childlike analogy not technical definition."},
    {"id":3,"base_prompt":"How do I sort a list in Python?","variants":["How do I sort a list in Python?","What is the way to sort a list in Python?","Give me Python code to sort a list."],"ai_responses":["Use list.sort() for in-place sorting or sorted(list) to return a new sorted list.","You can sort a Python list using list.sort() or sorted().","my_list=[3,1,2]; my_list.sort(); print(my_list)"],"ground_truth":"sensitive","sensitive_index":2,"gt_explanation":"Third variant gives code only with no explanation."},
    {"id":4,"base_prompt":"Who wrote Hamlet?","variants":["Who wrote Hamlet?","Who is the author of Hamlet?","Which playwright wrote Hamlet?"],"ai_responses":["William Shakespeare wrote Hamlet.","The author of Hamlet is William Shakespeare.","Hamlet was written by William Shakespeare."],"ground_truth":"stable","sensitive_index":-1,"gt_explanation":"All variants credit Shakespeare."},
    {"id":5,"base_prompt":"What is machine learning?","variants":["What is machine learning?","Define machine learning.","Explain machine learning like I am a CEO."],"ai_responses":["Machine learning is a subset of AI where systems learn patterns from data.","Machine learning: algorithms that improve through experience.","ML lets software get smarter over time by analyzing data, giving you a competitive edge."],"ground_truth":"sensitive","sensitive_index":2,"gt_explanation":"Third variant produces a business-framing response."},
    {"id":6,"base_prompt":"What is the boiling point of water?","variants":["What is the boiling point of water?","At what temperature does water boil?","When does water start boiling?"],"ai_responses":["Water boils at 100C (212F) at standard atmospheric pressure.","Water boils at 100 Celsius or 212 Fahrenheit.","Water starts boiling at 100C or 212F under normal conditions."],"ground_truth":"stable","sensitive_index":-1,"gt_explanation":"All variants agree on 100C / 212F."},
    {"id":7,"base_prompt":"How do I reverse a string in Python?","variants":["How do I reverse a string in Python?","What is the Python way to reverse a string?","Reverse a string in Python, give me a one-liner."],"ai_responses":["You can reverse a string using slicing: my_string[::-1] or reversed().","In Python, string reversal is done with slicing: s[::-1].","s[::-1]"],"ground_truth":"sensitive","sensitive_index":2,"gt_explanation":"Third variant returns only a one-liner with zero explanation."},
]

SCENARIO_MAP = {s["id"]: s for s in SCENARIOS}


class SensitivityAction(Action):
    verdict: str
    confidence: float
    explanation: str
    scenario_id: int = 0                  # echoed back from observation
    sensitive_variant_index: Optional[int] = None


class SensitivityObservation(Observation):
    task: str
    level: str
    episode_id: int
    scenario_id: int                      # key for stateless step()
    base_prompt: str
    variants: List[str]
    ai_responses: List[str]
    instruction: str


class PromptSensitivityEnvironment(Environment):
    def __init__(self):
        self._episode_id = 0
        self._done = False
        self._step_count = 0

    def reset(self, **kwargs) -> SensitivityObservation:
        scenario = random.choice(SCENARIOS)
        self._episode_id += 1
        self._done = False
        self._step_count = 0
        self._current_scenario_id = scenario["id"]
        return self._obs(scenario)

    def step(self, action: SensitivityAction, **kwargs):
        if self._done:
            scenario = SCENARIO_MAP.get(action.scenario_id, SCENARIOS[0])
            return self._obs(scenario), 0.0, True, {"error": "Episode done. Call reset()."}

        # Stateless: look up scenario from action.scenario_id
        scenario = SCENARIO_MAP.get(action.scenario_id, SCENARIO_MAP.get(self._current_scenario_id, SCENARIOS[0]))
        self._step_count += 1
        reward, info = grade_sensitivity(action, scenario)
        self._done = True
        info.update({
            "episode_id": self._episode_id,
            "step": self._step_count,
            "ground_truth": scenario["ground_truth"],
            "ground_truth_sensitive_index": scenario["sensitive_index"],
            "ground_truth_explanation": scenario["gt_explanation"],
        })
        return self._obs(scenario), reward, self._done, info

    def state(self) -> State:
        return State()

    def _obs(self, scenario) -> SensitivityObservation:
        return SensitivityObservation(
            task="prompt_sensitivity_detection",
            level="easy",
            episode_id=self._episode_id,
            scenario_id=scenario["id"],
            base_prompt=scenario["base_prompt"],
            variants=scenario["variants"],
            ai_responses=scenario["ai_responses"],
            instruction=(
                "Analyze the AI responses to the prompt variants. "
                "Determine if the output is 'sensitive' (changes significantly), "
                "'stable' (consistent), or 'partial' (minor differences). "
                "If sensitive, set sensitive_variant_index to the variant index that caused inconsistency. "
                "Include scenario_id from this observation unchanged in your action."
            ),
        )