env.py

"""
ETL Pipeline Fixer — Environment Implementation
================================================
Phase 3: Multi-level task logic.

Difficulty levels:
  - easy:   Schema mismatch (fix INSERT column name)
  - medium: Silent data loss (change INNER JOIN to LEFT JOIN)
  - hard:   Concurrency (fix 'database is locked' via threading.Lock)
"""

from __future__ import annotations

import shutil
import sqlite3
import subprocess
import sys
import textwrap
import uuid
from pathlib import Path
from typing import Any, Literal, Optional

from openenv.core.env_server.interfaces import Environment
from openenv.core.env_server.types import State

try:
    from .models import (
        EditCodeAction,
        EtlAction,
        EtlObservation,
        EtlReward,
        QueryDbAction,
        RunPipelineAction,
    )
except ImportError:
    from models import (
        EditCodeAction,
        EtlAction,
        EtlObservation,
        EtlReward,
        QueryDbAction,
        RunPipelineAction,
    )

# ---------------------------------------------------------------------------
# ── Types ──────────────────────────────────────────────────────────────────
# ---------------------------------------------------------------------------

DifficultyLevel = Literal["easy", "medium", "hard"]

# ---------------------------------------------------------------------------
# ── Task 1 (Easy) constants ────────────────────────────────────────────────
# ---------------------------------------------------------------------------

_EASY_EXPECTED_ROWS: list[tuple[int, str]] = [
    (1, "Alice Smith"),
    (2, "Bob Jones"),
    (3, "Charlie Brown"),
]

_EASY_SOURCE_CSV: str = textwrap.dedent("""\
    user_id,full_name
    1,Alice Smith
    2,Bob Jones
    3,Charlie Brown
""")

_EASY_PIPELINE: str = textwrap.dedent("""\
    \"\"\"ETL pipeline — Task 1: load source_data.csv into the users table.\"\"\"
    import csv
    import sqlite3

    # --- Extract ---
    rows = []
    with open("source_data.csv", newline="") as f:
        reader = csv.DictReader(f)
        for row in reader:
            rows.append(row)

    # --- Load ---
    conn = sqlite3.connect("etl_target.db")
    cursor = conn.cursor()

    # BUG: 'user_id' does not exist in the target schema — should be 'id'.
    for row in rows:
        cursor.execute(
            "INSERT INTO users (user_id, name) VALUES (?, ?)",
            (row["user_id"], row["full_name"]),
        )

    conn.commit()
    conn.close()

    print(f"Loaded {len(rows)} rows into users table.")
""")

# ---------------------------------------------------------------------------
# ── Task 2 (Medium) constants ──────────────────────────────────────────────
# ---------------------------------------------------------------------------

_MEDIUM_PIPELINE: str = textwrap.dedent("""\
    \"\"\"ETL pipeline — Task 2: Create user_summary from users and purchases.\"\"\"
    import sqlite3

    conn = sqlite3.connect("etl_target.db")
    cursor = conn.cursor()

    # BUG: INNER JOIN drops users without purchases.
    # The fix: change to LEFT JOIN to retain all 5 users.
    query = '''
        CREATE TABLE user_summary AS
        SELECT u.id, u.name, COUNT(p.id) as purchase_count
        FROM users u
        INNER JOIN purchases p ON u.id = p.user_id
        GROUP BY u.id, u.name
    '''
    cursor.execute("DROP TABLE IF EXISTS user_summary")
    cursor.execute(query)
    
    conn.commit()
    conn.close()
    print("User summary created.")
""")

# ---------------------------------------------------------------------------
# ── Task 3 (Hard) constants ────────────────────────────────────────────────
# ---------------------------------------------------------------------------

_HARD_PIPELINE: str = textwrap.dedent("""\
    \"\"\"ETL pipeline — Task 3: Concurrent ingestion.\"\"\"
    import sqlite3
    import threading

    # Shared connection (check_same_thread=False allows sharing)
    conn = sqlite3.connect("etl_target.db", check_same_thread=False)

    def worker(worker_id):
        cursor = conn.cursor()
        for i in range(20):
            # BUG: Threads share a single SQLite connection without a lock
            # leading to 'database is locked' errors or dropped inserts.
            # FIX: import threading and use a threading.Lock() around execute/commit
            try:
                cursor.execute(
                    "INSERT INTO events (worker_id, event_no) VALUES (?, ?)", 
                    (worker_id, i)
                )
                conn.commit()
            except sqlite3.OperationalError as e:
                print(f"Error in worker {worker_id}: {e}")

    threads = []
    for i in range(5):
        t = threading.Thread(target=worker, args=(i,))
        threads.append(t)
        t.start()

    for t in threads:
        t.join()
        
    conn.close()
    print("Data ingestion complete.")
""")



# ---------------------------------------------------------------------------
# ── Environment ──────────────────────────────────────────────────────────────
# ---------------------------------------------------------------------------

class ETLEnv(Environment):
    SUPPORTS_CONCURRENT_SESSIONS: bool = True
    _DEFAULT_TIMEOUT_S: float = 30.0
    _EDITABLE_FILES: frozenset[str] = frozenset({"pipeline.py"})

    def __init__(self, difficulty: DifficultyLevel = "easy", **kwargs: Any) -> None:
        super().__init__(**kwargs)
        self.difficulty: DifficultyLevel = difficulty
        self._workspace: Path = Path(__file__).parent / "workspaces" / str(uuid.uuid4())
        self._script_content: str = ""
        self._console_buffer: str = ""
        self._db_summary: str = ""
        self._done: bool = False
        self._grader_reward: float = 0.0
        self._state: State = State(episode_id=str(uuid.uuid4()), step_count=0)

    def reset(
        self,
        seed: Optional[int] = None,
        episode_id: Optional[str] = None,
        difficulty: Optional[DifficultyLevel] = None,
        **kwargs: Any,
    ) -> EtlObservation:
        requested_difficulty = difficulty or kwargs.get("difficulty")
        if requested_difficulty is not None:
            if requested_difficulty not in ("easy", "medium", "hard"):
                raise ValueError("difficulty must be one of: easy, medium, hard")
            self.difficulty = requested_difficulty

        self._reset_rubric()
        self._state = State(episode_id=episode_id or str(uuid.uuid4()), step_count=0)
        self._done = False
        self._grader_reward = 0.0
        self._console_buffer = ""

        if self._workspace.exists():
            shutil.rmtree(self._workspace)
        self._workspace.mkdir(parents=True, exist_ok=True)

        db_path = self._workspace / "etl_target.db"
        pipeline_path = self._workspace / "pipeline.py"
        
        conn = sqlite3.connect(str(db_path))

        # Setup varies by difficulty
        if self.difficulty == "easy":
            (self._workspace / "source_data.csv").write_text(_EASY_SOURCE_CSV, encoding="utf-8")
            conn.execute("CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT NOT NULL)")
            conn.commit()
            pipeline_path.write_text(_EASY_PIPELINE, encoding="utf-8")
            self._script_content = _EASY_PIPELINE

        elif self.difficulty == "medium":
            conn.execute("CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT NOT NULL)")
            conn.execute("CREATE TABLE purchases (id INTEGER PRIMARY KEY, user_id INTEGER, item TEXT NOT NULL)")
            # 5 users total
            for i, name in enumerate(["Alice", "Bob", "Charlie", "David", "Eve"], 1):
                conn.execute("INSERT INTO users (id, name) VALUES (?, ?)", (i, name))
            # 3 purchases for users 1, 2, 3 only
            conn.execute("INSERT INTO purchases (user_id, item) VALUES (1, 'Book')")
            conn.execute("INSERT INTO purchases (user_id, item) VALUES (2, 'Pen')")
            conn.execute("INSERT INTO purchases (user_id, item) VALUES (3, 'Laptop')")
            conn.commit()
            pipeline_path.write_text(_MEDIUM_PIPELINE, encoding="utf-8")
            self._script_content = _MEDIUM_PIPELINE

        elif self.difficulty == "hard":
            conn.execute("CREATE TABLE events (id INTEGER PRIMARY KEY, worker_id INTEGER, event_no INTEGER)")
            conn.commit()
            pipeline_path.write_text(_HARD_PIPELINE, encoding="utf-8")
            self._script_content = _HARD_PIPELINE

        conn.close()
        self._db_summary = self._summarize_db()

        return EtlObservation(
            console_logs=self._console_buffer,
            current_script_content=self._script_content,
            db_state_summary=self._db_summary,
            done=self._done,
            reward=0.0,
        )

    def step(
        self,
        action: EtlAction,  # type: ignore[override]
        timeout_s: Optional[float] = None,
        **kwargs: Any,
    ) -> EtlObservation:
        action_type = getattr(action, "action_type", None)

        if action_type == "edit_code":
            if isinstance(action, EditCodeAction):
                edit_action = action
            else:
                edit_action = EditCodeAction(
                    filepath=action.filepath,
                    target_line=action.target_line,
                    new_code=action.new_code,
                )
            self._apply_edit(edit_action)
        elif action_type == "run_pipeline":
            if isinstance(action, RunPipelineAction):
                run_action = action
            else:
                run_action = RunPipelineAction(timeout_override_s=action.timeout_override_s)
            self._execute_run(run_action)
        elif action_type == "query_db":
            if isinstance(action, QueryDbAction):
                query_action = action
            else:
                query_action = QueryDbAction(
                    sql_string=action.sql_string,
                    max_rows=action.max_rows,
                )
            self._execute_query(query_action)
        else:
            raise TypeError(f"Unrecognised action type: {type(action).__name__!r}")

        self._state.step_count += 1

        partial_obs = EtlObservation(
            console_logs=self._console_buffer,
            current_script_content=self._script_content,
            db_state_summary=self._db_summary,
            done=self._done,
            reward=0.0,
        )

        if self.rubric is not None:
            raw_reward = self._apply_rubric(action, partial_obs)
        else:
            raw_reward = self._grader_reward

        reward_obj = EtlReward(value=max(0.0, min(1.0, float(raw_reward))))

        return EtlObservation(
            console_logs=self._console_buffer,
            current_script_content=self._script_content,
            db_state_summary=self._db_summary,
            done=self._done,
            reward=reward_obj.value,
        )

    @property
    def state(self) -> State:
        return self._state

    def close(self) -> None:
        if self._workspace.exists():
            shutil.rmtree(self._workspace, ignore_errors=True)

    def _apply_edit(self, action: EditCodeAction) -> None:
        requested_basename = Path(action.filepath).name
        if requested_basename not in self._EDITABLE_FILES:
            self._console_buffer = f"[EDIT ERROR] '{action.filepath}' is not allowed."
            return

        target_path = (self._workspace / requested_basename).resolve()
        if not str(target_path).startswith(str(self._workspace.resolve())):
            self._console_buffer = f"[EDIT ERROR] Path traversal detected."
            return

        if not target_path.exists():
            self._console_buffer = f"[EDIT ERROR] File not found."
            return

        lines = target_path.read_text(encoding="utf-8").splitlines(keepends=True)

        if action.target_line > len(lines):
            self._console_buffer = f"[EDIT ERROR] target_line {action.target_line} out of range."
            return

        replacement = action.new_code
        if not replacement.endswith("\n"):
            replacement += "\n"
        lines[action.target_line - 1] = replacement

        updated_text = "".join(lines)
        target_path.write_text(updated_text, encoding="utf-8")

        self._script_content = updated_text
        self._console_buffer = f"[EDIT OK] Line {action.target_line} replaced."
        self._grader_reward = 0.0

    def _execute_run(self, action: RunPipelineAction) -> None:
        pipeline_path = self._workspace / "pipeline.py"
        if not pipeline_path.exists():
            self._console_buffer = "[RUN ERROR] pipeline.py missing."
            self._grader_reward = 0.0
            return

        timeout = action.timeout_override_s or self._DEFAULT_TIMEOUT_S

        try:
            result = subprocess.run(
                [sys.executable, str(pipeline_path)],
                capture_output=True,
                text=True,
                timeout=timeout,
                cwd=str(self._workspace),
            )
        except subprocess.TimeoutExpired:
            self._console_buffer = f"[RUN TIMEOUT] Exceeded {timeout:.1f}s."
            self._grader_reward = 0.0
            return
        except OSError as exc:
            self._console_buffer = f"[RUN OS ERROR] {exc}"
            self._grader_reward = 0.0
            return

        output_parts = []
        if result.stdout:
            output_parts.append(result.stdout.rstrip())
        if result.stderr:
            output_parts.append(f"[STDERR]\n{result.stderr.rstrip()}")

        exit_label = f"[Exit code: {result.returncode}]" if result.returncode != 0 else "[Exit code: 0 — OK]"
        output_parts.append(exit_label)
        self._console_buffer = "\n".join(output_parts)

        self._db_summary = self._summarize_db()

        # Score the task
        self._grader_reward = self._score_task()
        self._done = self._grader_reward >= 1.0

    def _execute_query(self, action: QueryDbAction) -> None:
        sql = action.sql_string.strip()
        if not sql.upper().startswith("SELECT"):
            self._db_summary = "[QUERY ERROR] Only SELECT allowed."
            self._grader_reward = 0.0
            return

        db_path = self._workspace / "etl_target.db"
        if not db_path.exists():
            self._db_summary = "[QUERY ERROR] DB missing."
            self._grader_reward = 0.0
            return

        try:
            conn = sqlite3.connect(str(db_path))
            cursor = conn.execute(sql)
            rows = cursor.fetchmany(action.max_rows)
            col_names = [d[0] for d in cursor.description] if cursor.description else []
            conn.close()
        except sqlite3.Error as exc:
            self._db_summary = f"[SQL ERROR] {exc}"
            self._grader_reward = 0.0
            return

        if not rows:
            self._db_summary = f"Query OK — 0 rows.\nCols: {col_names if col_names else '(none)'}"
        else:
            header = " | ".join(col_names)
            separator = "-" * max(len(header), 10)
            data_lines = [" | ".join(str(v) for v in row) for row in rows]
            self._db_summary = "\n".join([header, separator] + data_lines + [f"\n({len(rows)} rows)"])

        self._grader_reward = 0.0

    def _summarize_db(self) -> str:
        db_path = self._workspace / "etl_target.db"
        if not db_path.exists():
            return "Database missing."

        try:
            conn = sqlite3.connect(str(db_path))
            summary_parts = []
            if self.difficulty == "easy":
                (row_count,) = conn.execute("SELECT COUNT(*) FROM users").fetchone()
                sample = conn.execute("SELECT id, name FROM users LIMIT 3").fetchall()
                summary_parts.append(f"users: {row_count} rows | sample: {sample}")
            elif self.difficulty == "medium":
                (u_count,) = conn.execute("SELECT COUNT(*) FROM users").fetchone()
                (p_count,) = conn.execute("SELECT COUNT(*) FROM purchases").fetchone()
                summary_parts.append(f"users: {u_count} rows, purchases: {p_count} rows")
                # Check for user_summary existence safely
                cursor = conn.cursor()
                cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='user_summary'")
                if cursor.fetchone():
                    (us_count,) = conn.execute("SELECT COUNT(*) FROM user_summary").fetchone()
                    summary_parts.append(f"user_summary: {us_count} rows")
                else:
                    summary_parts.append("user_summary: not created yet")
            elif self.difficulty == "hard":
                (e_count,) = conn.execute("SELECT COUNT(*) FROM events").fetchone()
                summary_parts.append(f"events: {e_count} rows")
            conn.close()
            return "\n".join(summary_parts)
        except sqlite3.Error as exc:
            return f"[DB ERROR] {exc}"

    def _score_task(self) -> float:
        db_path = self._workspace / "etl_target.db"
        if not db_path.exists():
            return 0.0

        try:
            conn = sqlite3.connect(str(db_path))
            reward = 0.0
            
            if self.difficulty == "easy":
                try:
                    rows = conn.execute("SELECT id, name FROM users ORDER BY id").fetchall()
                    matches = sum(1 for expected in _EASY_EXPECTED_ROWS if expected in rows)
                    reward = min(matches / len(_EASY_EXPECTED_ROWS), 1.0)
                except sqlite3.Error:
                    pass

            elif self.difficulty == "medium":
                try:
                    cursor = conn.cursor()
                    cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='user_summary'")
                    if cursor.fetchone():
                        (count,) = conn.execute("SELECT COUNT(*) FROM user_summary").fetchone()
                        reward = min(max(count, 0) / 5.0, 1.0)
                except sqlite3.Error:
                    pass

            elif self.difficulty == "hard":
                try:
                    (count,) = conn.execute("SELECT COUNT(*) FROM events").fetchone()
                    reward = min(count / 100.0, 1.0)
                except sqlite3.Error:
                    pass

            conn.close()
            return reward
        except sqlite3.Error:
            return 0.0