Blame for server/environment.py

"""SQL Quality Environment - Core environment manager."""

from __future__ import annotations

from datetime import datetime
from difflib import get_close_matches
from itertools import combinations
import random as _random
import re
import sqlite3
import threading
import uuid
from typing import Any, Dict, List, Tuple

from fastapi import HTTPException

from .models import (
    ActionType,
    SQLAction,
    SQLObservation,
    SQLReward,
    SQLState,
    StepResponse,
    TrajectoryStep,
)
from .tasks import (
    create_dedup_bug_task,
    create_fk_integrity_task,
    create_null_audit_task,
    create_schema_drift_task,
    create_type_repair_task,
)

TASK_CREATORS = {
    "null_audit": create_null_audit_task,
    "type_repair": create_type_repair_task,
    "dedup_bug": create_dedup_bug_task,
    "fk_integrity": create_fk_integrity_task,
    "schema_drift": create_schema_drift_task,
}

NULLISH_TEXT_VALUES = {
    "",
    "n/a",
    "na",
    "nan",
    "null",
    "none",
    "missing",
    "unknown",
    "-1",
}
ALLOWED_READ_KEYWORDS = {"SELECT", "WITH", "EXPLAIN", "PRAGMA", "SHOW"}
IDENTITY_NAME_TOKENS = {
    "id",
    "title",
    "name",
    "email",
    "company",
    "location",
    "address",
    "original",
    "imdb",
}
BUSINESS_NUMERIC_TOKENS = {
    "score",
    "rating",
    "votes",
    "income",
    "revenue",
    "salary",
    "amount",
    "price",
    "cost",
    "duration",
    "founded",
}
LOW_SIGNAL_DEDUP_TOKENS = {
    "genre",
    "industry",
    "sector",
    "country",
    "release",
    "year",
    "type",
    "category",
}


def strip_sql_comments(sql: str) -> str:
    """Remove SQL comments while preserving string literals."""
    if not sql:
        return ""
    
    result = []
    i = 0
    n = len(sql)
    
    while i < n:
        char = sql[i]
        
        # Handle string literals
        if char in ("'", '"'):
            quote_char = char
            result.append(char)
            i += 1
            
            while i < n:
                char = sql[i]
                result.append(char)
                
                if char == quote_char:
                    # Check for doubled quote escape
                    if i + 1 < n and sql[i + 1] == quote_char:
                        result.append(sql[i + 1])  # Add the escaped quote
                        i += 2
                    else:
                        i += 1
                        break
                elif char == "\\" and i + 1 < n:
                    # Handle backslash escapes
                    result.append(sql[i + 1])
                    i += 2
                else:
                    i += 1
        elif char == "-" and i + 1 < n and sql[i + 1] == "-":
            # Skip -- comments until end of line
            i += 2
            while i < n and sql[i] != "\n":
                i += 1
        elif char == "/" and i + 1 < n and sql[i + 1] == "*":
            # Skip /* */ comments
            i += 2
            while i + 1 < n and not (sql[i] == "*" and sql[i + 1] == "/"):
                i += 1
            i += 2  # Skip the closing */
        else:
            result.append(char)
            i += 1
    
    return "".join(result).strip()


def split_sql_statements(sql: str) -> list[str]:
    """Split SQL into top-level statements without breaking string literals."""
    if not sql:
        return []
    
    statements: list[str] = []
    current = []
    in_string = False
    string_char = ""

    i = 0
    n = len(sql)
    
    while i < n:
        char = sql[i]
        prev = sql[i - 1] if i > 0 else ""
        next_char = sql[i + 1] if i < n - 1 else ""
        
        # Always append quote characters to current buffer
        if char in ("'", '"'):
            current.append(char)
            
            # Check for doubled-quote escape (SQL standard)
            if next_char == char:
                current.append(next_char)  # Add the escaped quote
                i += 1  # Skip the second quote
            elif not in_string:
                in_string = True
                string_char = char
            elif string_char == char:
                in_string = False
                string_char = ""
        elif char == "\\" and in_string and next_char:
            # Handle backslash escapes in strings
            current.append(char)
            current.append(next_char)
            i += 1
        elif char == ";" and not in_string:
            # Statement delimiter outside of strings
            statement = "".join(current).strip()
            if statement:
                statements.append(statement)
            current = []
        else:
            current.append(char)
        
        i += 1

    # Add the final statement if non-empty
    tail = "".join(current).strip()
    if tail:
        statements.append(tail)
    
    return statements


def first_sql_keyword(sql: str) -> str:
    """Return the first SQL keyword after stripping comments and whitespace."""
    cleaned = strip_sql_comments(sql)
    if not cleaned:
        return ""
    return cleaned.split()[0].upper()


def is_readonly_statement(sql: str) -> bool:
    """Return True when a statement is allowed in Zone 1 read-only execution."""
    # Remove comments and split into individual statements
    statements = split_sql_statements(sql)
    
    # Empty or whitespace-only SQL should return False
    if not statements:
        return False
    
    # Check every statement's first keyword
    for statement in statements:
        statement = statement.strip()
        if statement:  # Skip empty statements
            first_keyword = first_sql_keyword(statement)
            if not first_keyword or first_keyword not in ALLOWED_READ_KEYWORDS:
                return False
    
    return True


def quote_identifier(name: str) -> str:
    """Quote SQLite identifiers safely."""
    return '"' + str(name).replace('"', '""') + '"'


def collect_schema_identifiers(schema: Any) -> tuple[set[str], set[str], dict[str, str]]:
    """Collect table/column identifiers plus normalized lookup keys from schema-like objects."""
    tables: set[str] = set()
    columns: set[str] = set()
    normalized_lookup: dict[str, str] = {}

    for table_info in schema or []:
        if isinstance(table_info, dict):
            table_name = table_info.get("table_name", "")
            if table_name:
                tables.add(table_name)
                normalized_lookup.setdefault(normalize_schema_label(table_name), table_name)

            # Handle both string and dict column formats
            for col in table_info.get("columns", []):
                if isinstance(col, str):
                    column_name = col.split(" (")[0]
                else:
                    column_name = col.get("name", str(col))
                if column_name:
                    columns.add(column_name)
                    normalized_lookup.setdefault(normalize_schema_label(column_name), column_name)

    return tables, columns, normalized_lookup


def rewrite_known_quoted_identifiers(sql: str, schema: Any) -> Tuple[str, List[Tuple[str, str]], List[Tuple[str, str]]]:
    """Map quoted original CSV headers to sanitized schema identifiers when the match is exact after normalization."""
    sql_clean = strip_sql_comments(sql)
    tables, columns, normalized_lookup = collect_schema_identifiers(schema)
    alias_tokens = set(re.findall(r'(?i)\bAS\s+"([^"]+)"', sql_clean))
    rewrites: list[tuple[str, str]] = []
    unresolved: list[tuple[str, str]] = []
    rewritten_sql = sql

    for token in re.findall(r'"([^"]+)"', sql_clean):
        if token in alias_tokens or token in tables or token in columns:
            continue
        normalized = normalize_schema_label(token)
        if not normalized:
            continue
        replacement = normalized_lookup.get(normalized)
        if replacement and replacement != token:
            rewrites.append((token, replacement))
        else:
            close = get_close_matches(normalized, list(normalized_lookup.keys()), n=1, cutoff=0.75)
            if close:
                unresolved.append((token, normalized_lookup[close[0]]))

    for source, target in rewrites:
        # Use regex to match double-quoted identifiers only, not inside string literals
        pattern = fr'"({re.escape(source)})"'
        replacement = f'"{target}"'
        rewritten_sql = re.sub(pattern, replacement, rewritten_sql)

    # Deduplicate while preserving order.
    dedup_rewrites = list(dict.fromkeys(rewrites))
    dedup_unresolved = list(dict.fromkeys(unresolved))
    return rewritten_sql, dedup_rewrites, dedup_unresolved


def normalize_text(value: Any) -> str:
    """Normalize values for placeholder, boolean, and date heuristics."""
    if value is None:
        return ""
    return str(value).strip()


def normalize_schema_label(value: str) -> str:
    """Collapse punctuation and spacing for schema-aware comparisons."""
    import re
    return re.sub(r"[^a-z0-9]+", "", (value or "").lower())


def identifier_tokens(name: str) -> list[str]:
    """Split a column name into normalized semantic tokens."""
    import re
    return [token for token in re.split(r"[^a-z0-9]+", (name or "").lower()) if token]


def is_identifier_like_name(name: str) -> bool:
    """Return True for classic identifier columns such as *_id."""
    tokens = identifier_tokens(name)
    return bool(name and ((name or "").lower() == "id" or (tokens and tokens[-1] == "id")))


def is_identity_like_name(name: str) -> bool:
    """Return True for columns likely to participate in business keys."""
    tokens = set(identifier_tokens(name))
    return bool(tokens & IDENTITY_NAME_TOKENS)


def is_low_signal_dedup_name(name: str) -> bool:
    """Return True for categorical/date-ish names that rarely form a good dedup key alone."""
    tokens = set(identifier_tokens(name))
    return bool(tokens & LOW_SIGNAL_DEDUP_TOKENS)


def is_nullish_value(value: Any) -> bool:
    """Treat NULLs, empty strings, and common placeholders as missing."""
    if value is None:
        return True
    text = normalize_text(value).lower()
    return text in NULLISH_TEXT_VALUES


def clean_numeric_text(value: Any) -> str:
    """Remove common numeric noise such as currency and commas."""
    text = normalize_text(value)
    return re.sub(r"\s+", "", text.replace(",", "").replace("$", "").replace("%", ""))


def parse_numeric_scalar(value: Any) -> float | None:
    """Parse scalar numeric values while rejecting IDs and textual ranges."""
    text = clean_numeric_text(value)
    if not text:
        return None
    if re.fullmatch(r"[+-]?\d+(?:\.\d+)?", text):
        try:
            return float(text)
        except ValueError:
            return None
    suffix_match = re.fullmatch(r"([+-]?\d+(?:\.\d+)?)([KMB])", text, flags=re.IGNORECASE)
    if suffix_match:
        magnitude = {"K": 1_000, "M": 1_000_000, "B": 1_000_000_000}[suffix_match.group(2).upper()]
        try:
            return float(suffix_match.group(1)) * magnitude
        except ValueError:
            return None
    return None


def identifier_like_ratio(values: list[Any]) -> float:
    """Estimate how many values look like alphanumeric identifiers rather than measures."""
    clean_values = [normalize_text(v) for v in values if not is_nullish_value(v)]
    if not clean_values:
        return 0.0
    matches = 0
    for value in clean_values:
        compact = re.sub(r"[\s_-]+", "", value)
        if re.fullmatch(r"[A-Za-z]{1,8}\d{3,}", compact):
            matches += 1
    return matches / max(1, len(clean_values))


def numeric_parse_ratio(values: list[Any]) -> float:
    """Estimate how many non-nullish values can be parsed as numbers."""
    clean_values = [v for v in values if not is_nullish_value(v)]
    if not clean_values:
        return 0.0
    parsed = 0
    for value in clean_values:
        if parse_numeric_scalar(value) is not None:
            parsed += 1
    return parsed / max(1, len(clean_values))


def boolean_parse_ratio(values: list[Any]) -> float:
    """Estimate how many non-nullish values look like booleans."""
    clean_values = [normalize_text(v).lower() for v in values if not is_nullish_value(v)]
    if not clean_values:
        return 0.0
    boolean_tokens = {"true", "false", "yes", "no", "y", "n", "0", "1"}
    parsed = sum(1 for value in clean_values if value in boolean_tokens)
    return parsed / max(1, len(clean_values))


def date_parse_ratio(values: list[Any]) -> float:
    """Estimate how many non-nullish values look like dates or timestamps."""
    clean_values = [normalize_text(v) for v in values if not is_nullish_value(v)]
    if not clean_values:
        return 0.0
    parsed = 0
    patterns = [
        ("%Y-%m-%d",),
        ("%Y-%m-%d %H:%M:%S",),
        ("%Y-%m-%dT%H:%M:%S",),
        ("%m/%d/%Y",),
        ("%d/%m/%Y",),
    ]
    for value in clean_values:
        candidate = value.replace("Z", "")
        matched = False
        for (fmt,) in patterns:
            try:
                datetime.strptime(candidate, fmt)
                parsed += 1
                matched = True
                break
            except ValueError:
                continue
        if matched:
            continue
        try:
            datetime.fromisoformat(candidate)
            parsed += 1
        except ValueError:
            pass
    return parsed / max(1, len(clean_values))


def normalize_base_table_name(name: str) -> str:
    """Normalize versioned table names to support schema-drift detection."""
    text = name.lower()
    text = re.sub(r"(?:_v\d+|_version\d+|_ver\d+)$", "", text)
    text = re.sub(r"(?:\d{4}|\d+)$", "", text)
    return text.rstrip("_")


def fetch_sample_values(conn: sqlite3.Connection, table_name: str, column_name: str, limit: int = 200) -> list[Any]:
    """Fetch a bounded sample for heuristic analysis."""
    query = (
        f"SELECT {quote_identifier(column_name)} "
        f"FROM {quote_identifier(table_name)} "
        f"LIMIT {limit}"
    )
    return [row[0] for row in conn.execute(query).fetchall()]


def nullish_count_sql(table_name: str, column_name: str) -> str:
    """Build a deterministic SQL predicate for nullish values."""
    col = quote_identifier(column_name)
    return (
        f"SELECT COUNT(*) FROM {quote_identifier(table_name)} WHERE "
        f"{col} IS NULL OR TRIM(CAST({col} AS TEXT)) = '' OR "
        f"LOWER(TRIM(CAST({col} AS TEXT))) IN ('n/a','na','nan','null','none','missing','unknown','-1')"
    )


def analyze_uploaded_database(conn: sqlite3.Connection, schema_info: dict[str, dict]) -> dict[str, Any]:
    """Profile uploaded data and derive supported task contexts deterministically."""
    table_profiles: dict[str, dict[str, Any]] = {}
    task_candidates: dict[str, dict[str, Any]] = {}
    tables = list(schema_info.keys())

    for table_name in tables:
        table_schema = schema_info.get(table_name, {})
        raw_columns = table_schema.get("columns", [])
        columns = [
            col.split(" (")[0] if isinstance(col, str) else col.get("name", "")
            for col in raw_columns
        ]
        row_count = int(table_schema.get("row_count", 0) or 0)
        column_profiles: dict[str, dict[str, Any]] = {}
        duplicate_candidates: list[dict[str, Any]] = []

        for column_name in columns:
            try:
                nullish_count = conn.execute(nullish_count_sql(table_name, column_name)).fetchone()[0]
            except sqlite3.Error:
                nullish_count = 0
            try:
                header_like_count = conn.execute(
                    f"SELECT COUNT(*) FROM {quote_identifier(table_name)} "
                    f"WHERE LOWER(TRIM(CAST({quote_identifier(column_name)} AS TEXT))) = ?",
                    (column_name.lower(),),
                ).fetchone()[0]
            except sqlite3.Error:
                header_like_count = 0
            try:
                distinct_count = conn.execute(
                    f"SELECT COUNT(DISTINCT {quote_identifier(column_name)}) "
                    f"FROM {quote_identifier(table_name)}"
                ).fetchone()[0]
            except sqlite3.Error:
                distinct_count = 0

            sample_values = fetch_sample_values(conn, table_name, column_name)
            numeric_ratio = numeric_parse_ratio(sample_values)
            boolean_ratio = boolean_parse_ratio(sample_values)
            date_ratio = date_parse_ratio(sample_values)
            id_like_ratio = identifier_like_ratio(sample_values)
            numeric_samples = [parsed for parsed in (parse_numeric_scalar(v) for v in sample_values) if parsed is not None]
            semantic_anomaly = ""
            anomaly_ratio = 0.0
            name_tokens = set(identifier_tokens(column_name))
            if numeric_samples and name_tokens & {"score", "rating"}:
                too_large = sum(1 for value in numeric_samples if value > 10)
                in_range = sum(1 for value in numeric_samples if 0 <= value <= 10)
                if too_large > 0 and in_range >= max(3, int(len(numeric_samples) * 0.6)):
                    semantic_anomaly = "score_scale"
                    anomaly_ratio = round(too_large / max(1, len(numeric_samples)), 4)
            declared_type = ""
            for raw in raw_columns:
                if isinstance(raw, str) and raw.startswith(f"{column_name} "):
                    declared_type = raw.rsplit("(", 1)[-1].rstrip(")")
                    break

            column_profiles[column_name] = {
                "declared_type": declared_type.upper(),
                "nullish_count": nullish_count,
                "nullish_ratio": round(nullish_count / max(1, row_count), 4),
                "distinct_count": distinct_count,
                "distinct_ratio": round(distinct_count / max(1, row_count), 4),
                "numeric_ratio": round(numeric_ratio, 4),
                "boolean_ratio": round(boolean_ratio, 4),
                "date_ratio": round(date_ratio, 4),
                "identifier_like_ratio": round(id_like_ratio, 4),
                "semantic_anomaly": semantic_anomaly,
                "anomaly_ratio": anomaly_ratio,
                "header_like_count": header_like_count,
            }

            try:
                duplicate_groups = conn.execute(
                    f"SELECT COUNT(*) FROM ("
                    f"SELECT {quote_identifier(column_name)} "
                    f"FROM {quote_identifier(table_name)} "
                    f"WHERE {quote_identifier(column_name)} IS NOT NULL "
                    f"GROUP BY {quote_identifier(column_name)} "
                    f"HAVING COUNT(*) > 1)"
                ).fetchone()[0]
                expected_clean = conn.execute(
                    f"SELECT COUNT(DISTINCT {quote_identifier(column_name)}) "
                    f"FROM {quote_identifier(table_name)} "
                    f"WHERE {quote_identifier(column_name)} IS NOT NULL"
                ).fetchone()[0]
            except sqlite3.Error:
                duplicate_groups = 0
                expected_clean = row_count

            is_key_like = is_identifier_like_name(column_name) or column_name.lower() in {
                "id",
                "order_id",
                "invoice_id",
                "customer_id",
                "user_id",
                "event_id",
                "job_id",
            }
            distinct_ratio = column_profiles[column_name]["distinct_ratio"]
            if duplicate_groups >= 2 and (
                is_key_like or (is_identity_like_name(column_name) and distinct_ratio >= 0.7)
            ):
                duplicate_candidates.append(
                    {
                        "table": table_name,
                        "key_columns": [column_name],
                        "duplicate_groups": duplicate_groups,
                        "expected_clean_count": expected_clean,
                        "score": (
                            distinct_ratio
                            + (0.45 if is_key_like else 0.0)
                            + (0.2 if is_identity_like_name(column_name) else 0.0)
                            - min(0.3, duplicate_groups / max(1, row_count))
                        ),
                    }
                )

        if not duplicate_candidates and len(columns) >= 2:
            composite_pool = [
                column_name
                for column_name in columns
                if is_identity_like_name(column_name)
                or column_profiles[column_name]["distinct_ratio"] >= 0.55
            ][:6]
            for combo_size in (3, 2):
                for combo in combinations(composite_pool, combo_size):
                    if not any(is_identity_like_name(column_name) for column_name in combo):
                        continue
                    if all(is_low_signal_dedup_name(column_name) for column_name in combo):
                        continue
                    try:
                        key_expr = ", ".join(quote_identifier(column_name) for column_name in combo)
                        duplicate_groups = conn.execute(
                            f"SELECT COUNT(*) FROM ("
                            f"SELECT {key_expr} "
                            f"FROM {quote_identifier(table_name)} "
                            f"GROUP BY {key_expr} "
                            f"HAVING COUNT(*) > 1)"
                        ).fetchone()[0]
                        expected_clean = conn.execute(
                            f"SELECT COUNT(*) FROM ("
                            f"SELECT DISTINCT {key_expr} "
                            f"FROM {quote_identifier(table_name)})"
                        ).fetchone()[0]
                    except sqlite3.Error:
                        duplicate_groups = 0
                        expected_clean = row_count
                    if duplicate_groups <= 1:
                        continue
                    combined_distinct_ratio = expected_clean / max(1, row_count)
                    if combined_distinct_ratio < 0.6:
                        continue
                    semantic_bonus = 0.12 * sum(
                        1 for column_name in combo if is_identity_like_name(column_name)
                    )
                    duplicate_candidates.append(
                        {
                            "table": table_name,
                            "key_columns": list(combo),
                            "duplicate_groups": duplicate_groups,
                            "expected_clean_count": expected_clean,
                            "score": combined_distinct_ratio + semantic_bonus,
                        }
                    )
                if duplicate_candidates:
                    break

        table_profiles[table_name] = {
            "table_name": table_name,
            "row_count": row_count,
            "columns": columns,
            "column_profiles": column_profiles,
            "duplicate_candidates": sorted(
                duplicate_candidates,
                key=lambda item: item["score"],
                reverse=True,
            ),
        }

    # null_audit candidate
    null_candidate = None
    for table_profile in table_profiles.values():
        for column_name, profile in table_profile["column_profiles"].items():
            if profile["nullish_ratio"] < 0.05:
                continue
            candidate = {
                "task_id": "null_audit",
                "supported": True,
                "confidence": min(0.99, profile["nullish_ratio"] + 0.2),
                "table": table_profile["table_name"],
                "related_tables": [table_profile["table_name"]],
                "column": column_name,
                "description": f"Identify the highest nullish-rate column in {table_profile['table_name']}.",
                "input_mode": "answer",
                "input_label": "Column name",
                "input_placeholder": column_name,
                "action_guidance": "Enter only the single column name with the highest nullish rate.",
                "context_hint": (
                    f"Target table: {table_profile['table_name']}; focus on column {column_name}. "
                    f"Nullish ratio: {round(profile['nullish_ratio'] * 100, 1)}%."
                ),
                "verification_sql": (
                    f"SELECT COUNT(*) AS nullish_count FROM {quote_identifier(table_profile['table_name'])} "
                    f"WHERE {quote_identifier(column_name)} IS NULL "
                    f"OR TRIM(CAST({quote_identifier(column_name)} AS TEXT)) = ''"
                ),
                "ground_truth": {
                    "task_type": "null_audit",
                    "table": table_profile["table_name"],
                    "column": column_name,
                    "null_column": column_name,
                    "null_rate": profile["nullish_ratio"],
                    "original_row_count": table_profile["row_count"],
                    "all_columns": table_profile["columns"],
                    "uploaded_mode": True,
                },
            }
            if null_candidate is None or candidate["ground_truth"]["null_rate"] > null_candidate["ground_truth"]["null_rate"]:
                null_candidate = candidate
    task_candidates["null_audit"] = null_candidate or {
        "task_id": "null_audit",
        "supported": False,
        "reason": "No column with meaningful nullish values was detected in the uploaded data.",
    }

    # type_repair candidate
    type_candidate = None
    for table_profile in table_profiles.values():
        for column_name, profile in table_profile["column_profiles"].items():
            declared = profile["declared_type"]
            family = None
            ratio = 0.0
            name_tokens = set(identifier_tokens(column_name))
            if is_identifier_like_name(column_name) or profile.get("identifier_like_ratio", 0.0) >= 0.5:
                continue
            anomaly_type = profile.get("semantic_anomaly", "")
            if anomaly_type == "score_scale":
                family = "score_scale"
                ratio = max(0.8, 1.0 - profile.get("anomaly_ratio", 0.0))
            elif declared not in {"INTEGER", "REAL", "NUMERIC", "FLOAT", "DOUBLE"} and profile["numeric_ratio"] >= 0.75:
                family = "numeric"
                ratio = profile["numeric_ratio"]
            elif declared not in {"BOOLEAN", "INTEGER"} and profile["boolean_ratio"] >= 0.85:
                family = "boolean"
                ratio = profile["boolean_ratio"]
            elif declared not in {"DATE", "DATETIME", "TIMESTAMP"} and profile["date_ratio"] >= 0.75:
                family = "date"
                ratio = profile["date_ratio"]
            if not family:
                continue
            if family == "date" and ratio >= 0.98:
                # Fully parseable date-like text does not need repair unless another anomaly exists.
                continue
            score = ratio
            if family in {"numeric", "score_scale"}:
                score += 0.15
            if name_tokens & BUSINESS_NUMERIC_TOKENS:
                score += 0.2
            if name_tokens & {"score", "rating"} and family == "score_scale":
                score += 0.25
            if declared in {"", "TEXT"}:
                score += 0.05
            candidate = {
                "task_id": "type_repair",
                "supported": True,
                "confidence": min(0.99, score),
                "ranking_score": score,
                "table": table_profile["table_name"],
                "related_tables": [table_profile["table_name"]],
                "column": column_name,
                "description": (
                    f"Repair numeric anomalies in {table_profile['table_name']}.{column_name}."
                    if family == "score_scale"
                    else f"Repair typed values stored inconsistently in {table_profile['table_name']}.{column_name}."
                ),
                "input_mode": "fix_sql",
                "input_label": "Fix SQL",
                "input_placeholder": f'UPDATE {quote_identifier(table_profile["table_name"])} SET {quote_identifier(column_name)} = ...',
                "action_guidance": "Submit SQL that repairs the affected values on a cloned database.",
                "context_hint": (
                    f"Target table: {table_profile['table_name']}; column {column_name} contains score/rating values above 10 that likely need rescaling."
                    if family == "score_scale"
                    else f"Target table: {table_profile['table_name']}; column {column_name} looks {family}-like "
                    f"but is declared as {declared or 'TEXT'}."
                ),
                "verification_sql": (
                    f"SELECT {quote_identifier(column_name)}, typeof({quote_identifier(column_name)}) "
                    f"FROM {quote_identifier(table_profile['table_name'])} "
                    f"WHERE {quote_identifier(column_name)} IS NOT NULL LIMIT 10"
                ),
                "ground_truth": {
                    "task_type": "type_repair",
                    "table": table_profile["table_name"],
                    "column": column_name,
                    "target_column": column_name,
                    "issue_type": family,
                    "baseline_parse_ratio": ratio,
                    "baseline_anomaly_ratio": profile.get("anomaly_ratio", 0.0),
                    "declared_type": declared or "TEXT",
                    "original_row_count": table_profile["row_count"],
                    "all_columns": table_profile["columns"],
                    "uploaded_mode": True,
                },
            }
            if type_candidate is None or candidate["ranking_score"] > type_candidate.get("ranking_score", type_candidate["confidence"]):
                type_candidate = candidate
    task_candidates["type_repair"] = type_candidate or {
        "task_id": "type_repair",
        "supported": False,
        "reason": "No strongly typed-as-text column was detected in the uploaded data.",
    }

    # dedup_bug candidate
    dedup_candidate = None
    for table_profile in table_profiles.values():
        if not table_profile["duplicate_candidates"]:
            continue
        best = table_profile["duplicate_candidates"][0]
        key_columns = best["key_columns"]
        candidate = {
            "task_id": "dedup_bug",
            "supported": True,
            "confidence": 0.85,
            "table": table_profile["table_name"],
            "related_tables": [table_profile["table_name"]],
            "column": key_columns[0],
            "key_columns": key_columns,
            "description": f"Repair duplicate rows in {table_profile['table_name']}.",
            "input_mode": "fix_sql",
            "input_label": "Fix SQL",
            "input_placeholder": f'DELETE FROM {quote_identifier(table_profile["table_name"])} WHERE rowid NOT IN (...)',
            "action_guidance": "Submit SQL that removes duplicate rows while preserving one valid row per duplicate key.",
            "context_hint": f"Target table: {table_profile['table_name']}; duplicate key candidate: {', '.join(best.get('key_columns', ['unknown']))}",
            "ranking_score": best["score"],
            "ground_truth": {
                "task_type": "dedup_bug",
                "table": table_profile["table_name"],
                "key_columns": key_columns,
                "column": key_columns[0] if key_columns else None,
                "expected_clean_count": table_profile["row_count"] - best.get("duplicate_count", 0),
                "original_row_count": table_profile["row_count"],
                "uploaded_mode": True,
            }
        }
        if dedup_candidate is None or best["score"] > dedup_candidate.get("ranking_score", 0):
            dedup_candidate = candidate
    task_candidates["dedup_bug"] = dedup_candidate or {
        "task_id": "dedup_bug",
        "supported": False,
        "reason": "No meaningful duplicate-key pattern was detected in the uploaded data.",
    }

    # fk_integrity candidate
    fk_candidate = None
    for child_table, child_profile in table_profiles.items():
        for column_name in child_profile["columns"]:
            if not column_name.lower().endswith("_id") or column_name.lower() == "id":
                continue
            stem = column_name[:-3].lower()
            for parent_table, parent_profile in table_profiles.items():
                if parent_table == child_table:
                    continue
                parent_keys = [col for col in parent_profile["columns"] if col.lower() == "id" or col.lower() == stem or col.lower() == f"{stem}_id"]
                if not parent_keys:
                    continue
                parent_key = parent_keys[0]
                try:
                    orphan_count = conn.execute(
                        f"SELECT COUNT(*) FROM {quote_identifier(child_table)} c "
                        f"LEFT JOIN {quote_identifier(parent_table)} p "
                        f"ON c.{quote_identifier(column_name)} = p.{quote_identifier(parent_key)} "
                        f"WHERE c.{quote_identifier(column_name)} IS NOT NULL "
                        f"AND p.{quote_identifier(parent_key)} IS NULL"
                    ).fetchone()[0]
                    valid_count = conn.execute(
                        f"SELECT COUNT(*) FROM {quote_identifier(child_table)} c "
                        f"JOIN {quote_identifier(parent_table)} p "
                        f"ON c.{quote_identifier(column_name)} = p.{quote_identifier(parent_key)}"
                    ).fetchone()[0]
                except sqlite3.Error:
                    orphan_count = 0
                    valid_count = 0
                if orphan_count <= 0 or valid_count <= 0:
                    continue
                candidate = {
                    "task_id": "fk_integrity",
                    "supported": True,
                    "confidence": min(0.95, 0.5 + orphan_count / max(1, child_profile["row_count"])),
                    "table": child_table,
                    "related_tables": [child_table, parent_table],
                    "column": column_name,
                    "description": f"Resolve orphaned rows in {child_table}.{column_name}.",
                    "input_mode": "fix_sql",
                    "input_label": "Fix SQL",
                    "input_placeholder": f'UPDATE {quote_identifier(child_table)} SET {quote_identifier(column_name)} = ... WHERE ...',
                    "action_guidance": "Submit SQL that resolves orphan child rows without damaging valid relationships.",
                    "context_hint": (
                        f"Child table: {child_table}; parent table: {parent_table}; "
                        f"foreign key candidate: {column_name} -> {parent_key}."
                    ),
                    "verification_sql": (
                        f"SELECT COUNT(*) AS orphan_count FROM {quote_identifier(child_table)} c "
                        f"LEFT JOIN {quote_identifier(parent_table)} p "
                        f"ON c.{quote_identifier(column_name)} = p.{quote_identifier(parent_key)} "
                        f"WHERE c.{quote_identifier(column_name)} IS NOT NULL "
                        f"AND p.{quote_identifier(parent_key)} IS NULL"
                    ),
                    "ground_truth": {
                        "task_type": "fk_integrity",
                        "table": child_table,
                        "column": column_name,
                        "parent_table": parent_table,
                        "parent_key": parent_key,
                        "expected_clean_count": valid_count,
                        "orphan_count": orphan_count,
                        "original_row_count": child_profile["row_count"],
                        "all_columns": child_profile["columns"],
                        "uploaded_mode": True,
                    },
                }
                if fk_candidate is None or candidate["ground_truth"]["orphan_count"] > fk_candidate["ground_truth"]["orphan_count"]:
                    fk_candidate = candidate
    task_candidates["fk_integrity"] = fk_candidate or {
        "task_id": "fk_integrity",
        "supported": False,
        "reason": "No real parent-child orphan pattern was detected across the uploaded tables.",
    }

    # schema_drift candidate
    drift_candidate = None
    for left_table, right_table in combinations(tables, 2):
        left_base = normalize_base_table_name(left_table)
        right_base = normalize_base_table_name(right_table)
        left_cols = set(table_profiles[left_table]["columns"])
        right_cols = set(table_profiles[right_table]["columns"])
        overlap = len(left_cols & right_cols)
        if left_base != right_base and overlap < 2:
            continue
        if overlap == 0:
            continue
        canonical_name = "canonical_events" if "event" in left_base or "event" in right_base else f"canonical_{left_base or right_base}"
        expected_columns = sorted(left_cols | right_cols)
        candidate = {
            "task_id": "schema_drift",
            "supported": True,
            "confidence": min(0.9, overlap / max(1, min(len(left_cols), len(right_cols))) + 0.2),
            "table": canonical_name,
            "related_tables": [left_table, right_table],
            "column": expected_columns[0] if expected_columns else "",
            "description": f"Create a canonical table from {left_table} and {right_table}.",
            "input_mode": "fix_sql",
            "input_label": "Fix SQL",
            "input_placeholder": f'CREATE TABLE {quote_identifier(canonical_name)} AS SELECT ...',
            "action_guidance": "Submit SQL that creates a canonical table covering both source tables.",
            "context_hint": (
                f"Comparable source tables detected: {left_table} and {right_table}. "
                f"Canonical output should be {canonical_name}."
            ),
            "verification_sql": f"SELECT COUNT(*) FROM {quote_identifier(canonical_name)}",
            "ground_truth": {
                "task_type": "schema_drift",
                "table": canonical_name,
                "canonical_table": canonical_name,
                "source_tables": [left_table, right_table],
                "expected_columns": expected_columns,
                "expected_row_count": table_profiles[left_table]["row_count"] + table_profiles[right_table]["row_count"],
                "original_row_count": table_profiles[left_table]["row_count"] + table_profiles[right_table]["row_count"],
                "uploaded_mode": True,
            },
        }
        drift_candidate = candidate
        break
    task_candidates["schema_drift"] = drift_candidate or {
        "task_id": "schema_drift",
        "supported": False,
        "reason": "No comparable versioned tables were detected in the uploaded database.",
    }

    return {
        "table_profiles": table_profiles,
        "task_candidates": task_candidates,
    }


class SQLQualityEnvironment:
    """Main environment class for SQL Quality tasks."""

    def __init__(self) -> None:
        """Initialize environment."""
        self._current_task_id: str | None = None
        self._episode_id: str | None = None
        self._attempt_number: int = 1
        self._step_count: int = (
            0  # increments on every action (QUERY, PROFILE, FIX, SUBMIT)
        )
        self._max_steps_per_episode: int = 20  # prevent infinite loops
        self._cumulative_reward: float = 0.0
        self._done: bool = False
        self._ground_truth: Dict[str, Any] = {}
        self._cached_observation: SQLObservation | None = None
        self._lock = threading.Lock()
        # NEW: Multi-step trajectory support
        self._trajectory: list = []
        self._original_snapshot: dict = {}
        self._original_conn: sqlite3.Connection | None = None
        self._uploaded_analysis: dict[str, Any] = {}
        self._task_context: dict[str, Any] = {}
        self._active: bool = False
        self._last_feedback: str | None = None

    def reset(self, task_id: str, seed: int | None = None) -> SQLObservation:
        """Reset environment for a new episode.

        Args:
            task_id: One of the registered task IDs
            seed: Optional deterministic episode seed

        Returns:
            SQLObservation for the new episode

        Raises:
            ValueError: If task_id is not valid
        """
        with self._lock:
            # ── NEW: close any live connection from the previous episode ──────
            if hasattr(self, "_original_conn") and self._original_conn is not None:
                try:
                    self._original_conn.close()
                except Exception:
                    pass  # Already closed or never opened — safe to ignore
                self._original_conn = None

            # Validate task_id
            if task_id not in TASK_CREATORS:
                raise ValueError(f"Invalid task_id: {task_id}")

            # Generate new episode
            self._episode_id = str(uuid.uuid4())
            self._attempt_number = 1
            self._step_count = (
                0  # increments on every action (QUERY, PROFILE, FIX, SUBMIT)
            )
            self._cumulative_reward = 0.0
            self._done = False
            self._current_task_id = task_id
            self._trajectory = []
            self._original_snapshot = {}
            self._ground_truth = {}
            self._task_context = {}
            self._cached_observation = None
            self._last_feedback = None

            # Create task-specific database to extract observation data
            conn = None
            try:
                # Generate random episode seed for diversity
                episode_seed = seed if seed is not None else _random.randint(0, 2**31 - 1)
                creator = TASK_CREATORS[task_id]
                conn, observation, self._ground_truth = creator(episode_seed)

                # Store original connection for graders (read-only use)
                self._original_conn = conn

                # Snapshot the target table for data preservation grading
                target_table = self._ground_truth.get("table", "")
                if target_table:
                    from .graders import _snapshot_table

                    self._original_snapshot = _snapshot_table(conn, target_table)

                # Mark episode as active
                self._active = True

                # Cache observation for reuse in step()
                self._cached_observation = observation.model_copy(
                    update={"episode_id": self._episode_id}
                )
                
                # Store schema info for dynamic updates
                self._schema_info = observation.schema_info

                # Keep connection open for multi-step operations
                # Connection will be closed when episode ends or environment is reset

            except Exception as e:
                if conn is not None:
                    conn.close()
                raise e

            return self._cached_observation

    def reset_with_connection(
        self, 
        task_id: str, 
        conn: sqlite3.Connection,
        schema_info: dict,
        uploaded_analysis: dict[str, Any] | None = None,
    ) -> SQLObservation:
        """Reset environment using an uploaded database connection.
        
        This allows RL tasks to run on user-uploaded databases instead of
        synthetic data.

        Args:
            task_id: One of the registered task IDs
            conn: SQLite connection to the uploaded database
            schema_info: Dictionary mapping table names to schema info

        Returns:
            SQLObservation for the new episode

        Raises:
            ValueError: If task_id is not valid
        """
        with self._lock:
            # ── close any live connection from the previous episode ──────
            if hasattr(self, "_original_conn") and self._original_conn is not None:
                try:
                    self._original_conn.close()
                except Exception:
                    pass
                self._original_conn = None

            # Validate task_id
            if task_id not in TASK_CREATORS:
                raise ValueError(f"Invalid task_id: {task_id}")

            # Generate new episode
            self._episode_id = str(uuid.uuid4())
            self._attempt_number = 1
            self._step_count = 0
            self._cumulative_reward = 0.0
            self._done = False
            self._current_task_id = task_id
            self._trajectory = []
            self._original_snapshot = {}
            self._ground_truth = {}
            self._task_context = {}
            self._cached_observation = None
            self._last_feedback = None

            try:
                # Use the uploaded connection
                self._original_conn = conn
                self._uploaded_analysis = uploaded_analysis or analyze_uploaded_database(conn, schema_info)
                task_candidates = self._uploaded_analysis.get("task_candidates", {})
                task_candidate = task_candidates.get(task_id, {})

                if not task_candidate.get("supported"):
                    reason = task_candidate.get(
                        "reason",
                        f"Task '{task_id}' is not supported for the uploaded database.",
                    )
                    raise ValueError(reason)

                self._ground_truth = task_candidate.get("ground_truth", {}).copy()
                self._task_context = {
                    "task_id": task_id,
                    "target_table": task_candidate.get("table"),
                    "target_column": task_candidate.get("column"),
                    "related_tables": task_candidate.get("related_tables", []),
                    "key_columns": task_candidate.get("key_columns", []),
                    "verification_sql": task_candidate.get("verification_sql"),
                    "context_hint": task_candidate.get("context_hint"),
                    "confidence": task_candidate.get("confidence"),
                    "input_mode": task_candidate.get("input_mode", "fix_sql"),
                    "input_label": task_candidate.get("input_label", "Fix SQL"),
                    "input_placeholder": task_candidate.get("input_placeholder"),
                    "action_guidance": task_candidate.get("action_guidance"),
                    "description": task_candidate.get("description", f"Investigate task {task_id} on uploaded data."),
                    "max_attempts": task_candidate.get("max_attempts", 1),
                    "investigation_hint": task_candidate.get("context_hint", f"Using uploaded database."),
                }
                target_table = self._ground_truth.get("table", "")

                # Snapshot the target table for data preservation grading
                if target_table:
                    from .graders import _snapshot_table
                    self._original_snapshot = _snapshot_table(conn, target_table)

                # Build observation from uploaded schema
                from .models import SchemaInfo, SQLObservation

                schema_list = []
                sample_rows = {}

                related_tables = task_candidate.get("related_tables") or ([target_table] if target_table else [])
                for related_table in related_tables:
                    if related_table not in schema_info:
                        continue
                    table_data = schema_info[related_table]
                    schema_list.append(SchemaInfo(
                        table_name=related_table,
                        columns=table_data.get('columns', []),
                        row_count=table_data.get('row_count', 0)
                    ))

                    try:
                        cursor = conn.execute(f'SELECT * FROM "{related_table}" LIMIT 3')
                        cols = [d[0] for d in cursor.description]
                        rows = cursor.fetchall()
                        sample_rows[related_table] = [dict(zip(cols, row)) for row in rows]
                    except Exception:
                        sample_rows[related_table] = []

                observation = SQLObservation(
                    task_id=task_id,
                    task_description=task_candidate.get("description", f"Investigate task {task_id} on uploaded data."),
                    schema_info=schema_list,
                    sample_rows=sample_rows,
                    attempt_number=1,
                    episode_id=self._episode_id,
                    trajectory_history=[],
                    investigation_hint=task_candidate.get("context_hint", f"Using uploaded database: {target_table} table"),
                    max_attempts=1,
                )

                # Mark episode as active
                self._active = True
                self._cached_observation = observation
                self._schema_info = schema_list

            except Exception as e:
                raise e

            return self._cached_observation

    def task_context(self) -> dict[str, Any]:
        """Expose the current uploaded-data task context for UI helpers."""
        return self._task_context.copy()

    def _analyze_data_for_ground_truth(
        self, 
        conn: sqlite3.Connection, 
        task_id: str,
        target_table: str
    ) -> dict:
        """Analyze uploaded database to determine ground truth for scoring.
        
        Args:
            conn: SQLite connection
            task_id: The task type
            target_table: The primary table for this task
            
        Returns:
            Dictionary with ground truth information
        """
        ground_truth = {"task_type": task_id, "table": target_table}
        
        try:
            if task_id == "null_audit" and target_table:
                # Find column with highest null rate
                cursor = conn.execute(f'PRAGMA table_info("{target_table}")')
                columns = [row[1] for row in cursor.fetchall()]
                
                total_rows = conn.execute(
                    f'SELECT COUNT(*) FROM "{target_table}"'
                ).fetchone()[0]
                
                if total_rows > 0:
                    max_null_rate = 0
                    max_null_col = None
                    
                    for col in columns:
                        try:
                            null_count = conn.execute(
                                f'SELECT COUNT(*) FROM "{target_table}" WHERE "{col}" IS NULL'
                            ).fetchone()[0]
                            null_rate = null_count / total_rows
                            if null_rate > max_null_rate:
                                max_null_rate = null_rate
                                max_null_col = col
                        except Exception:
                            continue
                    
                    ground_truth["column"] = max_null_col
                    ground_truth["null_rate"] = max_null_rate
                    ground_truth["null_column"] = max_null_col
                    
            elif task_id == "dedup_bug" and target_table:
                # Find duplicate keys
                cursor = conn.execute(f'PRAGMA table_info("{target_table}")')
                columns = [row[1] for row in cursor.fetchall()]
                ground_truth["all_columns"] = columns
                
                if columns:
                    # Try to find duplicates on first column
                    key_col = columns[0]
                    ground_truth["column"] = key_col
                    ground_truth["offending_key"] = key_col
                    
                    # Count total duplicates
                    cursor = conn.execute(f'''
                        SELECT COUNT(*) FROM (
                            SELECT "{key_col}"
                            FROM "{target_table}" 
                            GROUP BY "{key_col}" 
                            HAVING COUNT(*) > 1
                        )
                    ''')
                    dup_groups = cursor.fetchone()[0]
                    
                    # Count expected rows after dedup (unique keys)
                    cursor = conn.execute(f'''
                        SELECT COUNT(DISTINCT "{key_col}") FROM "{target_table}"
                    ''')
                    expected_clean = cursor.fetchone()[0]
                    ground_truth["expected_clean_count"] = expected_clean
                    
                    # Get first duplicate example
                    cursor = conn.execute(f'''
                        SELECT "{key_col}", COUNT(*) as cnt 
                        FROM "{target_table}" 
                        GROUP BY "{key_col}" 
                        HAVING cnt > 1
                        LIMIT 1
                    ''')
                    result = cursor.fetchone()
                    if result:
                        ground_truth["duplicate_key"] = result[0]
                        ground_truth["duplicate_count"] = result[1]
                        ground_truth["duplicate_groups"] = dup_groups
                        
            elif task_id == "type_repair" and target_table:
                # Find columns that might have type issues (TEXT containing numbers)
                cursor = conn.execute(f'PRAGMA table_info("{target_table}")')
                columns = [row[1] for row in cursor.fetchall()]
                ground_truth["all_columns"] = columns
                
                # Find columns with numeric-like text values
                for col in columns:
                    try:
                        # Check if column contains numeric-looking values stored as TEXT
                        cursor = conn.execute(f'''
                            SELECT "{col}" FROM "{target_table}" 
                            WHERE "{col}" IS NOT NULL 
                            LIMIT 100
                        ''')
                        values = [row[0] for row in cursor.fetchall()]
                        
                        # Check if majority look like numbers but stored as text
                        numeric_like = 0
                        for v in values:
                            if v and str(v).strip():
                                try:
                                    float(str(v).replace(',', '').replace('$', ''))
                                    numeric_like += 1
                                except (ValueError, TypeError):
                                    pass
                        
                        if numeric_like > len(values) * 0.5:
                            ground_truth["column"] = col
                            ground_truth["target_column"] = col
                            ground_truth["issue_type"] = "TEXT_STORED_NUMERIC"
                            break
                    except Exception:
                        continue
                
                # If no type issue found, default to first column
                if "column" not in ground_truth and columns:
                    ground_truth["column"] = columns[0]
                    ground_truth["target_column"] = columns[0]
                    
            elif task_id == "fk_integrity" and target_table:
                # For FK integrity, we need to simulate orphan records
                # Since uploaded data may not have parent/child tables,
                # we'll create a virtual parent table based on existing data
                cursor = conn.execute(f'PRAGMA table_info("{target_table}")')
                columns = [row[1] for row in cursor.fetchall()]
                ground_truth["all_columns"] = columns
                
                # Look for columns that might be foreign keys (ending in _id)
                fk_candidates = [c for c in columns if c.lower().endswith('_id') and c.lower() != 'id']
                
                if fk_candidates:
                    fk_col = fk_candidates[0]
                    ground_truth["column"] = fk_col
                    
                    # Find values that would be orphans (using modulo to simulate missing parents)
                    cursor = conn.execute(f'''
                        SELECT DISTINCT "{fk_col}" FROM "{target_table}" 
                        WHERE "{fk_col}" IS NOT NULL
                    ''')
                    all_ids = [row[0] for row in cursor.fetchall()]
                    
                    if all_ids:
                        # Mark some IDs as invalid to simulate orphans
                        rng = _random.Random(42)  # Deterministic for testing
                        valid_ids = rng.sample(all_ids, max(1, len(all_ids) * 2 // 3))
                        valid_ids_set = set(valid_ids)
                        invalid_ids = [v for v in all_ids if v not in valid_ids_set]
                        
                        ground_truth["parent_table"] = f"{target_table}_parent"
                        ground_truth["parent_key"] = fk_col.replace('_id', '')
                        ground_truth["valid_invoice_ids"] = valid_ids
                        ground_truth["expected_clean_count"] = len([v for v in all_ids if v in valid_ids])
                        ground_truth["orphan_count"] = len(invalid_ids)
                        
            elif task_id == "schema_drift" and target_table:
                # Schema drift compares two versions of same table
                # For uploaded data, we'll treat it as canonical and create expected schema
                cursor = conn.execute(f'PRAGMA table_info("{target_table}")')
                columns_info = cursor.fetchall()
                columns = [row[1] for row in columns_info]
                
                ground_truth["table"] = target_table
                ground_truth["canonical_table"] = target_table
                ground_truth["expected_columns"] = columns
                
                # Get row count
                cursor = conn.execute(f'SELECT COUNT(*) FROM "{target_table}"')
                ground_truth["expected_row_count"] = cursor.fetchone()[0]
                
                # Get distinct IDs if available
                if 'id' in [c.lower() for c in columns]:
                    id_col = next(c for c in columns if c.lower() == 'id')
                    cursor = conn.execute(f'SELECT COUNT(DISTINCT "{id_col}") FROM "{target_table}"')
                    ground_truth["expected_distinct_ids"] = cursor.fetchone()[0]
                        
        except Exception as e:
            # Ground truth analysis failed, but we can still run the task
            ground_truth["analysis_error"] = str(e)
            
        return ground_truth

    def step(self, action: SQLAction) -> StepResponse:
        """
        Execute one agent action and return the resulting observation and reward.

        Routes by action_type:
            QUERY   - execute SELECT on episode conn, return result in observation
            PROFILE - compute column statistics, return in observation
            FIX     - execute on fresh_test_db(), score partially, continue episode
            SUBMIT  - call compute_episode_score(), end episode
        """
        with self._lock:
            if not self._active:
                raise HTTPException(
                    status_code=422, detail="No active episode. Call reset() first."
                )
            if self._done:
                raise HTTPException(
                    status_code=422,
                    detail="Episode is already done. Call reset() to start a new episode.",
                )

            # Check max steps guard
            if self._step_count >= self._max_steps_per_episode:
                raise HTTPException(
                    status_code=422,
                    detail="Max investigation steps reached. Call SUBMIT or reset().",
                )

            action = self._normalize_action(action)

            action_type = action.action_type

            # --- QUERY action ---
            if action_type == ActionType.QUERY:
                return self._handle_query(action)

            # --- PROFILE action ---
            elif action_type == ActionType.PROFILE:
                return self._handle_profile(action)

            # --- FIX action ---
            elif action_type == ActionType.FIX:
                return self._handle_fix(action)

            # --- SUBMIT action (or legacy single-step) ---
            else:
                return self._handle_submit(action)

    def _normalize_action(self, action: SQLAction) -> SQLAction:
        """Infer v2 action_type for backward-compatible payloads without action_type."""
        if "action_type" in action.model_fields_set:
            return action

        inferred = ActionType.SUBMIT
        if action.column_name:
            inferred = ActionType.PROFILE
        elif action.fix_query:
            inferred = ActionType.FIX
        elif action.sql:
            inferred = ActionType.QUERY if is_readonly_statement(action.sql) else ActionType.FIX
        elif action.answer or action.root_cause_table or action.offending_key:
            inferred = ActionType.SUBMIT

        return action.model_copy(update={"action_type": inferred})

    def _handle_query(self, action: SQLAction) -> StepResponse:
        """
        Handle a QUERY action.

        Executes the SQL on the episode connection (supports multi-statement).
        Returns a result preview in the observation. Does NOT end the episode.
        Awards small step_reward (0.05) for queries that target the right area.
        """
        # Minimal implementation to isolate the issue
        sql = (action.sql or "").strip()
        
        # Execute the query directly
        try:
            cursor = self._original_conn.execute(sql)
            if cursor.description:
                rows = cursor.fetchmany(5)
                cols = [d[0] for d in cursor.description]
                result_preview = {
                    "type": "table",
                    "columns": cols,
                    "rows": [dict(zip(cols, row)) for row in rows],
                    "total_rows": len(rows),
                }
            else:
                result_preview = f"Query executed successfully. Rows affected: {cursor.rowcount}"
        except Exception as e:
            result_preview = f"SQL Error: {e}"
        
        # Add to trajectory
        self._trajectory.append(
            TrajectoryStep(
                step_number=len(self._trajectory) + 1,
                action_type=ActionType.QUERY,
                sql_submitted=sql,
                result_preview=str(result_preview)[:400],
                step_reward=0.01,
            )
        )
        
        # Build minimal observation
        from .models import SQLObservation
        obs = SQLObservation(
            task_id=getattr(self, '_current_task_id', 'unknown'),
            episode_id=getattr(self, '_episode_id', 'unknown'),
            task_description="Test task",
            schema_info=getattr(self, '_schema_info', []),
            sample_rows={},
            attempt_number=getattr(self, '_attempt_number', 1),
            max_attempts=1,
            previous_feedback=None,
            trajectory_history=[],
            available_actions=["query", "profile", "fix", "submit"],
            investigation_hint=None,
        )
        
        return StepResponse(
            observation=obs,
            reward=SQLReward(
                score=self._cumulative_reward,
                step_reward=0.01,
                feedback=str(result_preview)[:400],
                dimensions={},
            ),
            done=False,
            info={"action_type": "query", "step": len(self._trajectory)},
        )

    def _handle_profile(self, action: SQLAction) -> StepResponse:
        """
        Handle a PROFILE action — statistical summary of a named column.

        Computes: null rate, distinct count, data type distribution, min/max.
        Does NOT end the episode. Awards step_reward based on column relevance.
        """
        col = (action.column_name or "").strip()
        result_preview = ""

        if not col:
            result_preview = "Error: column_name required for PROFILE action."
        else:
            # Try to find which table has this column
            try:
                tables = self._original_conn.execute(
                    "SELECT name FROM sqlite_master WHERE type='table'"
                ).fetchall()
                profile_results = []
                for (tbl,) in tables:
                    try:
                        cols = [
                            row[1]
                            for row in self._original_conn.execute(
                                f"PRAGMA table_info({tbl})"
                            ).fetchall()
                        ]
                        if col in cols:
                            total = self._original_conn.execute(
                                f"SELECT COUNT(*) FROM {quote_identifier(tbl)}"
                            ).fetchone()[0]
                            nulls = self._original_conn.execute(
                                nullish_count_sql(tbl, col)
                            ).fetchone()[0]
                            distinct = self._original_conn.execute(
                                f"SELECT COUNT(DISTINCT {quote_identifier(col)}) FROM {quote_identifier(tbl)}"
                            ).fetchone()[0]
                            profile_results.append(
                                f"{tbl}.{col}: total={total}, nulls={nulls} ({100*nulls//total if total else 0}%), distinct={distinct}"
                            )
                    except sqlite3.Error:
                        pass
                result_preview = (
                    " | ".join(profile_results)
                    if profile_results
                    else f"Column '{col}' not found."
                )
            except sqlite3.Error as e:
                result_preview = f"Profile error: {e}"

        # Determine step_reward based on relevance to ground truth
        ground_table = self._ground_truth.get("table", "").lower()
        ground_col = self._ground_truth.get(
            "column", self._ground_truth.get("null_column", "")
        ).lower()
        col_lower = col.lower()

        if col_lower == ground_col:
            step_reward = 0.08  # On-target: profiling exact relevant column
        elif ground_table and ground_table in result_preview.lower():
            step_reward = 0.05  # Near-target: profiled something from right table
        elif col_lower in result_preview.lower() and len(col_lower) > 2:
            step_reward = 0.01  # Off-target: profiled something real but irrelevant
        else:
            step_reward = 0.01  # Minimal: invalid column or junk probe

        self._trajectory.append(
            TrajectoryStep(
                step_number=len(self._trajectory) + 1,
                action_type=ActionType.PROFILE,
                sql_submitted=f"PROFILE({col})",
                result_preview=str(result_preview)[:200] if result_preview else "",
                step_reward=step_reward,
            )
        )

        self._cumulative_reward = min(
            1.0, round(self._cumulative_reward + step_reward, 4)
        )
        self._last_feedback = str(result_preview)[:200] if result_preview else ""
        self._step_count += 1
        obs = self._build_observation(result_preview=result_preview)

        return StepResponse(
            observation=obs,
            reward=SQLReward(
                score=self._cumulative_reward,
                step_reward=step_reward,
                feedback=str(result_preview)[:200] if result_preview else "",
            ),
            done=False,
            info={"action_type": "profile", "column_profiled": col},
        )

    def _handle_fix(self, action: SQLAction) -> StepResponse:
        """
        Handle a FIX action — execute write query on isolated copy for preview scoring.

        Executes on a _fresh_test_db() copy. Does NOT mutate episode state.
        Returns intermediate feedback on whether fix appears to work.
        Does NOT end episode — agent can follow up with verification query or SUBMIT.
        Awards intermediate reward based on quick fix assessment.
        """
        fix_sql = action.effective_sql or ""
        step_reward = 0.0
        feedback = ""
        rewrite_note = ""

        if fix_sql:
            if getattr(self, "_schema_info", None):
                # fix_sql, rewrites, unresolved = rewrite_known_quoted_identifiers(fix_sql, self._schema_info)
                fix_sql, rewrites, unresolved = fix_sql, [], []
                if unresolved:
                    feedback = (
                        "Fix preview blocked: unknown quoted identifier(s). Closest schema matches: "
                        + ", ".join(f'"{src}" -> "{dst}"' for src, dst in unresolved[:3])
                    )
                    fix_sql = ""
                if rewrites:
                    rewrite_note = " Auto-mapped: " + ", ".join(
                        f'"{src}" -> "{dst}"' for src, dst in rewrites[:3]
                    ) + "."
            from .graders import grade_fix_correctness

            if fix_sql:
                quick_score, severity = grade_fix_correctness(
                    fix_sql,
                    self._original_conn,
                    self._original_snapshot,
                    self._ground_truth,
                )
                step_reward = quick_score * 0.3  # Partial credit for fix attempt
                preview_label = "Answer preview" if self._ground_truth.get("task_type") == "null_audit" else "Fix preview"
                feedback = f"{preview_label}: {severity} ({quick_score:.2f}). Submit to finalize score.{rewrite_note}"
        else:
            feedback = "No SQL provided in fix action."

        self._trajectory.append(
            TrajectoryStep(
                step_number=len(self._trajectory) + 1,
                action_type=ActionType.FIX,
                sql_submitted=fix_sql,
                result_preview=feedback[:200],
                step_reward=step_reward,
            )
        )

        self._cumulative_reward = min(
            1.0, round(self._cumulative_reward + step_reward, 4)
        )
        self._last_feedback = feedback
        self._step_count += 1
        obs = self._build_observation(result_preview=feedback)

        return StepResponse(
            observation=obs,
            reward=SQLReward(
                score=self._cumulative_reward,
                step_reward=step_reward,
                feedback=feedback,
            ),
            done=False,
            info={
                "action_type": "fix",
                "severity_preview": severity if fix_sql else "N/A",
            },
        )

    def _handle_submit(self, action: SQLAction) -> StepResponse:
        """
        Handle a SUBMIT action — compute final composite score and end episode.

        Calls compute_episode_score() with the full trajectory.
        Sets done=True. This is the only action that ends the episode.
        """
        from .graders import compute_episode_score

        reward = compute_episode_score(
            trajectory=self._trajectory,
            ground_truth=self._ground_truth,
            original_conn=self._original_conn,
            original_snapshot=self._original_snapshot,
            submit_action=action,
        )

        self._done = True
        self._active = False
        self._cumulative_reward = min(
            1.0, round(self._cumulative_reward + reward.score, 4)
        )
        self._last_feedback = reward.feedback
        obs = self._build_observation(result_preview="")

        # Close connection when episode ends
        if self._original_conn:
            self._original_conn.close()
            self._original_conn = None

        return StepResponse(
            observation=obs,
            reward=reward,
            done=True,
            info={
                "action_type": "submit",
                "final_score": reward.score,
                "dimensions": reward.dimensions,
            }
        )

        if fix_sql:
            # quick_score, severity = grade_fix_correctness(
            #     fix_sql,
            #     self._original_conn,
            #     self._original_snapshot,
            #     self._ground_truth,
            # )
            quick_score, severity = 0.5, "UNKNOWN"
            step_reward = quick_score * 0.3  # Partial credit for fix attempt
            preview_label = "Answer preview" if self._ground_truth.get("task_type") == "null_audit" else "Fix preview"
            feedback = f"{preview_label}: {severity} ({quick_score:.2f}). Submit to finalize score.{rewrite_note}"
        else:
            feedback = "No SQL provided in fix action."

        self._trajectory.append(
            TrajectoryStep(
                step_number=len(self._trajectory) + 1,
                action_type=ActionType.FIX,
                sql_submitted=fix_sql,
                result_preview=feedback[:200],
                step_reward=step_reward,
            )
        )

        self._cumulative_reward = min(
            1.0, round(self._cumulative_reward + step_reward, 4)
        )
        self._last_feedback = feedback[:200] if feedback else ""
        self._step_count += 1

        # Build observation
        obs = self._build_observation(result_preview=feedback)

        return StepResponse(
            observation=obs,
            reward=SQLReward(
                score=self._cumulative_reward,
                step_reward=step_reward,
                feedback=feedback[:200] if feedback else "",
                dimensions={},
            ),
            done=False,
            info={"action_type": "fix", "step": len(self._trajectory)},
        )

    def _build_observation(self, result_preview: str = "") -> SQLObservation:
        """Build observation with current state and optional result preview."""
        # Format trajectory history for frontend
        trajectory_history = []
        for step in self._trajectory:
            trajectory_history.append({
                "step_number": step.step_number,
                "action_type": step.action_type.value,
                "sql_submitted": step.sql_submitted,
                "result_preview": step.result_preview,
                "step_reward": step.step_reward,
            })

        # Fallback to _cached_observation for synthetic tasks where _task_context may be empty
        task_description = self._task_context.get("description") or (self._cached_observation.task_description if self._cached_observation else "")
        max_attempts = self._task_context.get("max_attempts") or (self._cached_observation.max_attempts if self._cached_observation else 1)
        investigation_hint = self._task_context.get("investigation_hint") or (self._cached_observation.investigation_hint if self._cached_observation else None)
        
        # Check if current task has a grader (all tasks currently have graders)
        # This makes available_actions dynamic based on task grader status
        tasks_with_graders = {"null_audit", "type_repair", "dedup_bug", "fk_integrity", "schema_drift"}
        has_grader = self._current_task_id in tasks_with_graders
        
        available_actions = ["query", "profile", "fix"]
        if has_grader:
            available_actions.append("submit")
        
        return SQLObservation(
            task_id=self._current_task_id,
            episode_id=self._episode_id,
            task_description=task_description or "",
            schema_info=self._schema_info,
            sample_rows=self._cached_observation.sample_rows if self._cached_observation else {},
            attempt_number=self._attempt_number,
            max_attempts=max_attempts or 1,
            previous_feedback=self._last_feedback,
            trajectory_history=trajectory_history,
            available_actions=available_actions,
            investigation_hint=investigation_hint,
        )

    def state(self) -> SQLState:
        """Get current environment state.

        Returns:
            SQLState with current episode information

        Raises:
            ValueError: If no episode has been started
        """
        if self._episode_id is None:
            raise ValueError("No episode started. Call reset() first.")

        return SQLState(
            task_id=self._current_task_id,
            episode_id=self._episode_id,
            attempt_number=self._attempt_number,
            cumulative_reward=self._cumulative_reward,
            done=self._done,
        )