server/environment.py

"""SQL Query Debugger Environment Implementation."""

import sqlite3
from typing import Any, Optional

from .models import (
    SQLAction,
    SQLObservation,
    StepResult,
    ResetResult,
    StateResult,
)
from .tasks import TaskDefinition, get_task, SQLGrader, list_tasks


class SQLDebuggerEnvironment:
    """
    OpenEnv-compliant environment for SQL query debugging.

    Agents must diagnose and fix broken SQL queries to produce
    the expected results.
    """

    def __init__(self, task_id: str = "easy_syntax_fix"):
        """Initialize the environment with a specific task."""
        self.task_id = task_id
        self.task: Optional[TaskDefinition] = None
        self.grader: Optional[SQLGrader] = None

        # Episode state
        self.current_step: int = 0
        self.max_steps: int = 10
        self.done: bool = False
        self.total_reward: float = 0.0
        self.attempts_used: int = 0
        self.hints_used: int = 0
        self.last_error: Optional[str] = None
        self.last_query_result: Optional[str] = None
        self.best_score: float = 0.0
        self.history: list[dict[str, Any]] = []

    def reset(self, task_id: Optional[str] = None) -> ResetResult:
        """
        Reset the environment to initial state.

        Args:
            task_id: Optional task to switch to

        Returns:
            ResetResult with initial observation
        """
        if task_id:
            self.task_id = task_id

        self.task = get_task(self.task_id)
        self.grader = SQLGrader(self.task)
        self.grader.setup_database()

        self.current_step = 0
        self.max_steps = self.task.max_steps
        self.done = False
        self.total_reward = 0.0
        self.attempts_used = 0
        self.hints_used = 0
        self.last_error = None
        self.last_query_result = None
        self.best_score = 0.0
        self.history = []

        # Execute broken query to show initial error
        success, _, error = self.grader.execute_query(self.task.broken_query)
        if not success:
            self.last_error = error

        observation = self._make_observation()
        return ResetResult(observation=observation)

    def step(self, action: SQLAction) -> StepResult:
        """
        Execute an action and return the result.

        Args:
            action: The action to execute

        Returns:
            StepResult with observation, reward, done, and info
        """
        if self.done:
            return StepResult(
                observation=self._make_observation(),
                reward=0.0,
                done=True,
                info={"error": "Episode already finished"}
            )

        self.current_step += 1
        reward = 0.0
        info: dict[str, Any] = {}

        if action.action_type == "submit_fix":
            reward, info = self._handle_submit_fix(action.query)
        elif action.action_type == "execute_query":
            reward, info = self._handle_execute_query(action.query)
        elif action.action_type == "request_hint":
            reward, info = self._handle_request_hint()
        else:
            info["error"] = f"Unknown action type: {action.action_type}"
            reward = -0.05  # Small penalty for invalid action

        # Record history
        self.history.append({
            "step": self.current_step,
            "action": action.model_dump(),
            "reward": reward,
            "done": self.done
        })

        self.total_reward += reward

        # Check if max steps reached
        if self.current_step >= self.max_steps and not self.done:
            self.done = True
            info["termination_reason"] = "max_steps_reached"

        observation = self._make_observation()
        return StepResult(
            observation=observation,
            reward=reward,
            done=self.done,
            info=info
        )

    def state(self) -> StateResult:
        """Return current environment state."""
        return StateResult(
            task_id=self.task_id,
            current_step=self.current_step,
            max_steps=self.max_steps,
            done=self.done,
            total_reward=self.total_reward,
            attempts_used=self.attempts_used,
            hints_used=self.hints_used
        )

    def _make_observation(self) -> SQLObservation:
        """Create observation from current state."""
        sample_data = self._format_sample_data()

        return SQLObservation(
            task_id=self.task_id,
            task_description=self.task.description,
            schema_ddl=self.task.schema_ddl.strip(),
            sample_data=sample_data,
            broken_query=self.task.broken_query.strip(),
            error_message=self.last_error,
            expected_output_hint=self.task.expected_output_hint,
            attempts_remaining=self.max_steps - self.current_step,
            current_step=self.current_step,
            max_steps=self.max_steps,
            last_query_result=self.last_query_result,
            hints_used=self.hints_used,
            available_actions=["submit_fix", "execute_query", "request_hint"]
        )

    def _format_sample_data(self) -> str:
        """Format sample data as readable table."""
        if not self.grader or not self.grader.conn:
            return "Sample data not available"

        lines = []
        cursor = self.grader.conn.cursor()

        # Get all table names
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
        tables = [row[0] for row in cursor.fetchall()]

        for table in tables:
            cursor.execute(f"SELECT * FROM {table}")
            rows = cursor.fetchall()
            columns = [desc[0] for desc in cursor.description]

            lines.append(f"\n-- {table} --")
            lines.append(" | ".join(columns))
            lines.append("-" * 50)
            for row in rows:
                lines.append(" | ".join(str(v) for v in row))

        return "\n".join(lines)

    def _handle_submit_fix(self, query: Optional[str]) -> tuple[float, dict[str, Any]]:
        """Handle a fix submission."""
        info: dict[str, Any] = {}

        if not query:
            info["error"] = "No query provided"
            return -0.05, info

        self.attempts_used += 1

        # Grade the submission
        score, reason, partial_scores = self.grader.grade(query)
        info["score"] = score
        info["reason"] = reason
        info["partial_scores"] = partial_scores

        # Calculate reward based on improvement and absolute score
        reward = self._calculate_reward(score, partial_scores)

        if score >= 0.99:
            self.done = True
            info["success"] = True
            # Bonus for solving with fewer steps
            steps_remaining = self.max_steps - self.current_step
            efficiency_bonus = 0.1 * (steps_remaining / self.max_steps)
            reward += efficiency_bonus
            info["efficiency_bonus"] = efficiency_bonus
        else:
            info["success"] = False
            # Store last error for feedback
            success, results, error = self.grader.execute_query(query)
            if not success:
                self.last_error = error
            else:
                self.last_error = f"Query executed but results incorrect: {reason}"
                if results:
                    self.last_query_result = str(results[:5])  # Show first 5 rows

        # Update best score
        if score > self.best_score:
            self.best_score = score

        return reward, info

    def _handle_execute_query(self, query: Optional[str]) -> tuple[float, dict[str, Any]]:
        """Handle a test query execution (not a submission)."""
        info: dict[str, Any] = {}

        if not query:
            info["error"] = "No query provided"
            return -0.02, info

        success, results, error = self.grader.execute_query(query)

        if success:
            self.last_error = None
            self.last_query_result = str(results[:10]) if results else "No results"
            info["results"] = results[:10] if results else []
            info["row_count"] = len(results) if results else 0
            # Small positive reward for successful exploration
            return 0.01, info
        else:
            self.last_error = error
            self.last_query_result = None
            info["error"] = error
            return 0.0, info  # No penalty for exploration

    def _handle_request_hint(self) -> tuple[float, dict[str, Any]]:
        """Handle a hint request."""
        info: dict[str, Any] = {}

        if self.hints_used >= len(self.task.hints):
            info["hint"] = "No more hints available"
            return 0.0, info

        hint = self.task.hints[self.hints_used]
        self.hints_used += 1
        info["hint"] = hint
        info["hints_remaining"] = len(self.task.hints) - self.hints_used

        # Small penalty for using hints
        return -0.02, info

    def _calculate_reward(self, score: float, partial_scores: dict[str, float]) -> float:
        """
        Calculate reward based on score and progress.

        Reward shaping to encourage:
        - Syntactically valid queries
        - Getting closer to correct results
        - Improvement over previous attempts
        """
        # Base reward from score
        reward = score * 0.5

        # Improvement bonus
        if score > self.best_score:
            improvement = score - self.best_score
            reward += improvement * 0.3

        # Partial progress bonuses
        if partial_scores.get("syntax_valid", 0) > 0:
            reward += 0.05  # Bonus for valid syntax
        if partial_scores.get("row_count", 0) > 0:
            reward += 0.03  # Bonus for correct row count
        if partial_scores.get("data_correct", 0) > 0.25:
            reward += 0.05  # Bonus for partial data match

        # Penalty for repeated low scores
        if self.attempts_used > 3 and score < 0.3:
            reward -= 0.05

        return round(reward, 4)

    def close(self):
        """Clean up resources."""
        if self.grader:
            self.grader.cleanup()


# Factory function for creating environment instances
def create_environment(task_id: str = "easy_syntax_fix") -> SQLDebuggerEnvironment:
    """Create a new SQL Debugger environment instance."""
    return SQLDebuggerEnvironment(task_id=task_id)


def get_available_tasks() -> list[str]:
    """Get list of available task IDs."""
    return list_tasks()