"""
SQL Debug Environment — server-side environment logic.
Implements the OpenEnv Environment interface: reset(), step(), state property.
"""

import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import random
from uuid import uuid4
from typing import Optional
from models import Action, Observation, State
from tasks import TASKS, TASK_GROUPS, get_task
from graders import grade_submission

MAX_STEPS = 5


class SQLDebugEnvironment:
    """
    Real-world OpenEnv environment for SQL query debugging.

    An agent is given a broken SQL query and must submit corrected versions
    to earn rewards. The environment tracks attempts and computes shaped rewards.
    """

    def __init__(self):
        self._state: Optional[State] = None
        self._current_task_id: Optional[str] = None
        self._current_task: Optional[dict] = None
        self._submissions: list = []
        self._done: bool = False
        self._last_reward: float = 0.0
        self._last_result: str = ""
        self._task_group: Optional[str] = None

    def reset(
        self,
        task_group: Optional[str] = None,
        task_id: Optional[str] = None,
        seed: Optional[int] = None,
    ) -> Observation:
        """
        Initialize a new episode.

        Args:
            task_group: One of 'easy_syntax_fix', 'medium_logic_fix', 'hard_optimization_fix'
            task_id: Specific task id (overrides task_group)
            seed: Random seed for reproducibility
        """
        if seed is not None:
            random.seed(seed)

        if task_id and task_id in TASKS:
            self._current_task_id = task_id
            self._task_group = None
        elif task_id and task_id in TASK_GROUPS:
            self._current_task_id = random.choice(TASK_GROUPS[task_id])
            self._task_group = task_id
        elif task_group and task_group in TASK_GROUPS:
            self._current_task_id = random.choice(TASK_GROUPS[task_group])
            self._task_group = task_group
        else:

            self._current_task_id = random.choice(list(TASKS.keys()))
            self._task_group = None

        self._current_task = get_task(self._current_task_id)
        self._submissions = []
        self._done = False
        self._last_reward = 0.0
        self._last_result = ""

        episode_id = str(uuid4())
        self._state = State(
            episode_id=episode_id,
            step_count=0,
            task_id=self._task_group if self._task_group else self._current_task_id,
            max_steps=MAX_STEPS,
            cumulative_reward=0.0,
        )

        return self._build_observation()

    def step(self, action: Action) -> Observation:
        """
        Execute one fix attempt.

        Args:
            action: Agent's submitted SQL query + optional explanation

        Returns:
            Observation with reward, done flag, and feedback
        """
        if self._done:
            return self._build_observation(
                message="Episode already done. Call reset()."
            )

        if self._state is None:
            raise RuntimeError("Environment not initialized. Call reset() first.")

        submitted_query = action.query.strip()
        self._submissions.append(submitted_query)
        self._state.step_count += 1

        score, feedback, details = grade_submission(
            self._current_task_id, submitted_query
        )

        reward = self._compute_reward(score, self._state.step_count)
        self._last_reward = reward
        self._state.cumulative_reward += reward
        self._last_result = feedback

        if score == 1.0:
            self._done = True
        elif self._state.step_count >= MAX_STEPS:
            self._done = True

        return self._build_observation()

    def state(self) -> State:
        """Return current environment state metadata."""
        if self._state is None:
            return State(episode_id="uninitialized", step_count=0)
        return self._state

    def _compute_reward(self, score: float, step_num: int) -> float:
        """
        Shaped reward function with partial progress signals.

        - score=1.0: full reward, bonus for solving early
        - score=0.7: meaningful partial credit
        - score=0.5: small partial credit
        - score=0.3: tiny credit for at least running
        - score=0.0: penalty for wasted attempt (after first)
        """
        if score == 1.0:

            early_bonus = max(0.0, (MAX_STEPS - step_num) * 0.05)
            return round(1.0 + early_bonus, 4)
        elif score >= 0.7:
            return 0.5
        elif score >= 0.5:
            return 0.2
        elif score >= 0.3:
            return 0.1
        else:

            return -0.05 if step_num > 1 else 0.0

    def _build_observation(self, message: str = "") -> Observation:
        task = self._current_task
        attempts_remaining = MAX_STEPS - (self._state.step_count if self._state else 0)

        return Observation(
            schema_sql=task["schema_sql"].strip(),
            sample_data=task["sample_data"].strip(),
            buggy_query=task["buggy_query"].strip(),
            error_message=task["error_message"],
            task_description=task["task_description"],
            expected_columns=task["expected_columns"],
            attempts_remaining=max(0, attempts_remaining),
            last_attempt_result=self._last_result + (" " + message if message else ""),
            reward=self._last_reward,
            done=self._done,
            success=self._done and self._last_reward >= 1.0,
        )