server/environment.py

"""
SQL Migration Environment Server Implementation.

This is the core environment that wraps SQLite and exposes it via the OpenEnv
Environment interface. Each WebSocket session gets its own environment instance
with an isolated in-memory database.

Architecture Fixes Applied:
- A1: SELECT queries return actual data rows (not just "rows affected")
- A2: SQL execution timeout via progress handler (prevents infinite CTEs)
- A3: Dangerous SQL blacklist (ATTACH, DETACH, LOAD_EXTENSION, writable_schema)
- A4: Transaction awareness (respects BEGIN/COMMIT/ROLLBACK from agent)
- A5: Trajectory logging (full SQL history in metadata on episode end)
- A6: Per-task max_steps from seeds registry
"""

import re
import sqlite3
import uuid
import difflib
from typing import Any, Dict, List, Optional

# Support both in-repo and standalone imports
try:
    from openenv.core.env_server.interfaces import Environment
    from models import MigrationAction, MigrationObservation, MigrationState
    from server.grader import StateReconciler
except ImportError:
    from openenv.core.env_server.interfaces import Environment
    from ..models import MigrationAction, MigrationObservation, MigrationState
    from .grader import StateReconciler

# Import seeds (handle both import paths)
try:
    from .. import seeds
except ImportError:
    import seeds


# --- A3: Dangerous SQL Blacklist ---
_DANGEROUS_PATTERNS = re.compile(
    r"\b(ATTACH\s+DATABASE|DETACH\s+DATABASE|LOAD_EXTENSION)\b"
    r"|PRAGMA\s+writable_schema",
    re.IGNORECASE,
)

# --- A4: Transaction control keywords ---
_TX_BEGIN = re.compile(r"^\s*(BEGIN|BEGIN\s+TRANSACTION|BEGIN\s+DEFERRED|BEGIN\s+IMMEDIATE|BEGIN\s+EXCLUSIVE)\s*;?\s*$", re.IGNORECASE)
_TX_END = re.compile(r"^\s*(COMMIT|END|END\s+TRANSACTION|ROLLBACK)\s*;?\s*$", re.IGNORECASE)

# --- A2: Maximum SQLite operations before timeout ---
_MAX_OPS = 500_000  # ~5 seconds on typical hardware


# (Timeout handled via progress handler return value, no exception needed)


class DbMigrationEnvironment(Environment):
    """
    SQL Schema Migration Environment.

    An AI agent is dropped into a broken or schema-drifted SQLite database
    and must write SQL to migrate it to the target state without losing data.

    Each instance is isolated (per-WebSocket session) with its own :memory: database.
    """

    SUPPORTS_CONCURRENT_SESSIONS = True

    def __init__(self, task_name: str = "column-restructure"):
        """
        Initialize the migration environment.

        Args:
            task_name: One of the registered task names in seeds.TASKS
        """
        super().__init__()

        if task_name not in seeds.TASKS:
            raise ValueError(
                f"Unknown task: {task_name}. "
                f"Available: {list(seeds.TASKS.keys())}"
            )

        self.task_name = task_name
        self._task_config = seeds.TASKS[task_name]
        self._conn: Optional[sqlite3.Connection] = None
        self._reconciler: Optional[StateReconciler] = None
        self._step_count = 0
        self._trajectory: List[Dict[str, Any]] = []  # A5
        self._in_explicit_tx = False  # A4
        self._max_steps = self._task_config.get("max_steps", 20)  # A6
        self._state = MigrationState(
            task_name=task_name,
            migration_progress=0.0,
            max_steps=self._max_steps,  # A6
        )

    def _get_current_schema(self) -> str:
        """Get current database schema as DDL string, filtering internal tables."""
        if self._conn is None:
            return ""
        try:
            cursor = self._conn.execute(
                "SELECT sql FROM sqlite_master WHERE type='table' "
                "AND sql IS NOT NULL AND name NOT LIKE 'sqlite_%' ORDER BY name"
            )
            schemas = [row[0] for row in cursor.fetchall()]
            return ";\n\n".join(schemas) + ";" if schemas else ""
        except Exception:
            return ""

    def _generate_erd(self) -> str:
        """Generate a Mermaid.js erDiagram based on the current database structure."""
        if self._conn is None:
            return ""
        try:
            lines = ["erDiagram"]
            
            # 1. Get all tables
            cursor = self._conn.execute(
                "SELECT name FROM sqlite_master WHERE type='table' "
                "AND name NOT LIKE 'sqlite_%' ORDER BY name"
            )
            tables = [row[0] for row in cursor.fetchall()]
            
            relationships = []
            
            for table in tables:
                lines.append(f"    {table} {{")
                # 2. Get column info for each table
                cursor = self._conn.execute(f"PRAGMA table_info([{table}])")
                for col in cursor.fetchall():
                    # col[1]: name, col[2]: type, col[5]: pk
                    name = col[1]
                    dtype = col[2].replace(" ", "_")
                    is_pk = "PK" if col[5] else ""
                    lines.append(f"        {dtype} {name} {is_pk}")
                lines.append("    }")
                
                # 3. Get foreign keys for relationships
                cursor = self._conn.execute(f"PRAGMA foreign_key_list([{table}])")
                for fk in cursor.fetchall():
                    # fk[2]: to_table, fk[3]: from_col, fk[4]: to_col
                    to_table = fk[2]
                    relationships.append(f"    {table} ||--o{{ {to_table} : \"references\"")
            
            # Append unique relationships to avoid bloat
            lines.extend(list(set(relationships)))
            return "\n".join(lines)
        except Exception:
            return "erDiagram\n    ERROR { string info }"

    def _is_read_query(self, sql: str) -> bool:
        """Check if SQL is a read-only query (SELECT or certain PRAGMAs)."""
        stripped = sql.strip().upper()
        if stripped.startswith("SELECT") or stripped.startswith("WITH"):
            return True
        # PRAGMA table_info, foreign_key_list, etc. are read-only
        if stripped.startswith("PRAGMA") and "=" not in stripped:
            return True
        return False

    def _execute_with_timeout(self, sql: str) -> tuple:
        """
        Execute SQL with a progress-handler-based timeout.

        Returns: (cursor_or_None, error_string_or_None)
        """
        ops_count = [0]

        def _progress_callback():
            ops_count[0] += 1
            if ops_count[0] > _MAX_OPS:
                return 1  # Non-zero = abort
            return 0

        self._conn.set_progress_handler(_progress_callback, 1000)
        try:
            cursor = self._conn.execute(sql)
            return cursor, None, ops_count[0]
        except sqlite3.OperationalError as e:
            err_str = str(e).lower()
            if "interrupted" in err_str or ops_count[0] > _MAX_OPS:
                return None, "Error: Query exceeded execution time limit (possible infinite loop). Simplify your query.", ops_count[0]
            if "table" in err_str and "already exists" in err_str:
                return None, f"Schema Error: {e}. You must DROP the old table first if replacing it.", ops_count[0]
            if "has no column" in err_str:
                return None, f"Schema Error: {e}. Check table columns.", ops_count[0]
            return None, str(e), ops_count[0]
        except sqlite3.Warning as e:
            # Multi-statement fallback
            try:
                self._conn.executescript(sql)
                return None, None, ops_count[0]
            except Exception as script_e:
                return None, f"Error (Multi-Statement Fallback Failed): {script_e}. Original error: {e}", ops_count[0]
        except Exception as e:
            err_str = str(e).lower()
            if "values for" in err_str and "columns" in err_str:
                return None, f"Data Error: {e}. Ensure you are inserting the correct number of columns.", ops_count[0]
            return None, str(e), ops_count[0]
        finally:
            self._conn.set_progress_handler(None, 0)

    def _format_query_results(self, cursor) -> str:
        """Format SELECT query results as a readable table string."""
        try:
            rows = cursor.fetchall()
            if not rows:
                return "Query returned 0 rows."

            # Get column names
            col_names = [desc[0] for desc in cursor.description] if cursor.description else []

            # Cap at 50 rows
            truncated = len(rows) > 50
            display_rows = rows[:50]

            # Build output
            header = " | ".join(col_names) if col_names else "Results"
            lines = [header, "-" * len(header)]
            for row in display_rows:
                lines.append(" | ".join(str(v) for v in row))
            if truncated:
                lines.append(f"... ({len(rows) - 50} more rows truncated)")
            lines.append(f"({len(rows)} rows total)")

            return "\n".join(lines)
        except Exception:
            return "Query executed successfully."

    def reset(
        self,
        seed: Optional[int] = None,
        episode_id: Optional[str] = None,
        **kwargs: Any,
    ) -> MigrationObservation:
        """
        Reset the environment: create a fresh in-memory database and seed it.

        Args:
            seed: Unused (deterministic environment)
            episode_id: Optional episode identifier
            **kwargs: Additional reset parameters (including task_name override)

        Returns:
            Initial MigrationObservation with the broken schema and target DDL
        """
        # Allow task_name override via reset kwargs
        task_name = kwargs.get("task_name", self.task_name)
        if task_name != self.task_name and task_name in seeds.TASKS:
            self.task_name = task_name
            self._task_config = seeds.TASKS[task_name]
            self._max_steps = self._task_config.get("max_steps", 20)

        # Clean up previous connection
        if self._conn is not None:
            try:
                self._conn.close()
            except Exception:
                pass
            self._conn = None

        # Create fresh in-memory database
        self._conn = sqlite3.connect(":memory:")

        # Performance PRAGMAs for Docker I/O
        self._conn.execute("PRAGMA journal_mode = MEMORY")
        
        # CRITICAL: Enable foreign key enforcement
        self._conn.execute("PRAGMA foreign_keys = ON")

        # Seed source data
        self._task_config["seed_fn"](self._conn)

        # Initialize grader
        self._reconciler = StateReconciler(self.task_name)

        # Reset counters and trajectory
        self._step_count = 0
        self._trajectory = []  # A5
        self._in_explicit_tx = False  # A4
        self._state = MigrationState(
            episode_id=episode_id or str(uuid.uuid4()),
            step_count=0,
            task_name=self.task_name,
            migration_progress=0.0,
            max_steps=self._max_steps,  # A6
        )

        # Compute initial score and sync grader baseline
        initial_score = self._reconciler.score(self._conn)
        self._reconciler._last_score = initial_score  # Prevent inflated first-step reward
        self._state.migration_progress = initial_score

        current_ddl = self._get_current_schema()
        target_ddl = self._task_config["target_ddl"]
        diff = "\n".join(difflib.unified_diff(
            current_ddl.splitlines(),
            target_ddl.splitlines(),
            fromfile="current_schema",
            tofile="target_schema",
            lineterm=""
        ))

        return MigrationObservation(
            done=False,
            reward=0.0,
            current_schema_sql=current_ddl,
            target_schema_sql=target_ddl,
            last_execution_result="Environment initialized. Ready for migration.",
            step_number=0,
            migration_progress=initial_score,
            task_name=self.task_name,
            schema_diff=diff if diff else "Schemas match exactly.",
            erd_visualization=self._generate_erd(),
            metadata={"status": "ready"},
        )

    def step(
        self,
        action: MigrationAction,
        timeout_s: Optional[float] = None,
        **kwargs: Any,
    ) -> MigrationObservation:
        """
        Execute a SQL action against the database.

        Args:
            action: MigrationAction with sql_command, reasoning, and submit_final
            timeout_s: Unused (we use progress handler instead)
            **kwargs: Additional parameters

        Returns:
            MigrationObservation with execution result, updated schema, and reward
        """
        if not isinstance(action, MigrationAction):
            raise ValueError(f"Expected MigrationAction, got {type(action)}")

        if self._conn is None or self._reconciler is None:
            return MigrationObservation(
                done=True,
                reward=0.0,
                current_schema_sql="",
                target_schema_sql=self._task_config["target_ddl"],
                last_execution_result="Error: Environment not initialized. Call reset() first.",
                step_number=self._step_count,
                migration_progress=0.01,
                task_name=self.task_name,
                metadata={"error": "not_initialized"},
            )

        import time
        start_time = time.time()

        self._step_count += 1
        sql_command = action.sql_command.strip()

        # --- A3: Dangerous SQL Blacklist ---
        sql_lower = sql_command.lower()
        if re.search(r"pragma\s+foreign_keys\s*=\s*(off|0)", sql_lower):
            execution_result = (
                "Security Error: Disabling PRAGMA foreign_keys is strictly forbidden. "
                "Tip: To maintain integrity, perform your migration using temporary tables "
                "or deferred constraints instead of disabling enforcement."
            )
            action_error = "pragma_off_blocked"
        elif _DANGEROUS_PATTERNS.search(sql_command):
            execution_result = (
                "Error: This SQL command is not allowed for security reasons. "
                "ATTACH DATABASE, DETACH DATABASE, LOAD_EXTENSION, and "
                "PRAGMA writable_schema are blocked. "
                "Tip: Use standard DML (INSERT/UPDATE/DELETE) and DDL (CREATE/DROP) within a single database."
            )
            action_error = "blocked_command"
        else:
            # --- A4: Transaction Awareness ---
            execution_result = ""
            action_error = None

            if _TX_BEGIN.match(sql_command):
                # Agent wants to start a transaction
                try:
                    self._conn.execute("BEGIN")
                    self._in_explicit_tx = True
                    execution_result = "Success: Transaction started."
                except Exception as e:
                    execution_result = str(e)
                    action_error = str(e)
            elif _TX_END.match(sql_command):
                # Agent wants to commit or rollback
                try:
                    if sql_command.strip().upper().startswith("ROLLBACK"):
                        self._conn.rollback()
                        execution_result = "Success: Transaction rolled back."
                    else:
                        self._conn.commit()
                        execution_result = "Success: Transaction committed."
                    self._in_explicit_tx = False
                except Exception as e:
                    execution_result = str(e)
                    action_error = str(e)
                    self._in_explicit_tx = False
            else:
                # --- Normal SQL execution with timeout (A1, A2) ---
                cursor, error, q_ops = self._execute_with_timeout(sql_command)

                if error:
                    execution_result = error
                    action_error = error
                    # Rollback failed transaction
                    try:
                        if not self._in_explicit_tx:
                            self._conn.rollback()
                    except Exception:
                        pass
                else:
                    # --- A1: SELECT result passthrough ---
                    if self._is_read_query(sql_command):
                        execution_result = self._format_query_results(cursor)
                    else:
                        rows_affected = getattr(cursor, "rowcount", -1) if cursor else -1
                        execution_result = f"Success: Action executed. Rows affected: {rows_affected}"
                        # Try to auto-commit
                        if not self._in_explicit_tx:
                            try:
                                self._conn.commit()
                            except Exception:
                                pass

        # Compute scores
        current_score, step_reward = self._reconciler.compute_step_reward(self._conn)

        # Episode termination: submit_final, max steps, OR perfect score
        done = action.submit_final or self._step_count >= self._max_steps or current_score >= 0.99

        # --- A5: Trajectory logging ---
        self._trajectory.append({
            "step": self._step_count,
            "sql": action.sql_command,
            "reasoning": action.reasoning,
            "result": execution_result[:200],  # Truncate for storage
            "score": current_score,
            "reward": step_reward,
            "error": action_error,
        })

        # Update state
        self._state.step_count = self._step_count
        self._state.migration_progress = current_score

        # Build metadata with reasoning and debug info
        execution_ms = int((time.time() - start_time) * 1000)
        q_ops = q_ops if 'q_ops' in locals() else 0
        meta = {
            "reasoning": action.reasoning,
            "sql_executed": action.sql_command,
            "step": self._step_count,
            "execution_ms": execution_ms,
            "query_ops": q_ops,
        }
        if action_error:
            meta["error"] = action_error
        # Include full trajectory on episode end
        if done:
            meta["trajectory"] = self._trajectory

        current_ddl = self._get_current_schema()
        target_ddl = self._task_config["target_ddl"]
        diff = "\n".join(difflib.unified_diff(
            current_ddl.splitlines(),
            target_ddl.splitlines(),
            fromfile="current_schema",
            tofile="target_schema",
            lineterm=""
        ))

        return MigrationObservation(
            done=done,
            reward=step_reward,
            current_schema_sql=current_ddl,
            target_schema_sql=target_ddl,
            last_execution_result=execution_result,
            step_number=self._step_count,
            migration_progress=current_score,
            task_name=self.task_name,
            schema_diff=diff if diff else "Schemas match exactly.",
            erd_visualization=self._generate_erd(),
            metadata=meta,
        )

    def state(self) -> MigrationState:
        """Get current environment state."""
        return self._state

    def close(self) -> None:
        """Clean up resources."""
        if self._conn is not None:
            try:
                self._conn.close()
            except Exception:
                pass
            self._conn = None