Blame for server/tasks.py

"""Task definitions for SQL Quality Environment."""

from __future__ import annotations

import sqlite3
from typing import Any, Dict, Tuple

from .models import SchemaInfo, SQLObservation

# Allowed table names for validation
ALLOWED_TABLES = {
    "users",
    "transactions",
    "categories",
    "orders",
    "order_items",
    "shipments",
    "customers",
    "invoices",
    "events_v1",
    "events_v2",
    "canonical_events",
}


def _validate_table_name(table_name: str) -> str:
    """Validate table name against allowed set to prevent SQL injection."""
    if table_name not in ALLOWED_TABLES:
        raise ValueError(f"Invalid table name: {table_name}")
    return table_name


def _extract_sample_rows(
    conn: sqlite3.Connection, table_name: str, limit: int = 5
) -> list:
    """Extract sample rows with correct column names as dict keys."""
    validated_name = _validate_table_name(table_name)
    cursor = conn.execute(f"SELECT * FROM {validated_name} LIMIT ?", (limit,))
    columns = [description[0] for description in cursor.description]
    rows = cursor.fetchall()
    return [dict(zip(columns, row)) for row in rows]


def create_null_audit_task(
    seed: int = 42,
) -> Tuple[sqlite3.Connection, SQLObservation, Dict[str, Any]]:
    """Create null audit task with users table."""
    import random

    rng = random.Random(seed)
    conn = sqlite3.connect(":memory:", check_same_thread=False)
    cursor = conn.cursor()

    nullable_columns = ["email", "age", "city", "signup_date"]
    null_column = rng.choice(nullable_columns)

    cursor.execute(
        """
        CREATE TABLE users (
            id INTEGER PRIMARY KEY,
            name TEXT NOT NULL,
            email TEXT,
            age INTEGER,
            city TEXT,
            signup_date TEXT,
            last_login_at TEXT,
            referral_code TEXT
        )
    """
    )

    total_rows = 1000
    target_null_count = 520
    null_indices = set(rng.sample(range(total_rows), target_null_count))
    users_data = []

    for i in range(total_rows):
        name = f"user_{i}"
        email = f"user_{i}@example.com"
        age = rng.randint(18, 80)
        city = rng.choice(["New York", "Los Angeles", "Chicago", "Houston", "Phoenix"])
        signup_date = f"2023-{rng.randint(1, 12):02d}-{rng.randint(1, 28):02d}"
        last_login_at = (
            f"2024-{rng.randint(1, 12):02d}-{rng.randint(1, 28):02d}"
            f"T{rng.randint(0, 23):02d}:{rng.randint(0, 59):02d}:00"
        )
        referral_code = (
            f"REF_{rng.choice('ABCDEFGHJKLMNPQRSTUVWXYZ')}"
            f"{rng.randint(0, 9)}"
            f"{rng.choice('ABCDEFGHJKLMNPQRSTUVWXYZ')}"
        )

        if i in null_indices:
            if null_column == "email":
                email = None
            elif null_column == "age":
                age = None
            elif null_column == "city":
                city = None
            elif null_column == "signup_date":
                signup_date = None

        if rng.random() < 0.1 and null_column != "email":
            email = None
        if rng.random() < 0.1 and null_column != "age":
            age = None
        if rng.random() < 0.1 and null_column != "city":
            city = None
        if rng.random() < 0.1 and null_column != "signup_date":
            signup_date = None

        users_data.append(
            (i, name, email, age, city, signup_date, last_login_at, referral_code)
        )

    cursor.executemany(
        """
        INSERT INTO users (
            id, name, email, age, city, signup_date, last_login_at, referral_code
        )
        VALUES (?, ?, ?, ?, ?, ?, ?, ?)
    """,
        users_data,
    )

    # Verify table has data
    cursor.execute("SELECT COUNT(*) FROM users")
    count = cursor.fetchone()[0]
    assert count > 0, f"users table must have >0 rows, got {count}"

    cursor.execute("PRAGMA table_info(users)")
    pragma_rows = cursor.fetchall()
    columns = [f"{row[1]} ({row[2]})" for row in pragma_rows]
    all_columns = [row[1] for row in pragma_rows]

    sample_rows = {"users": _extract_sample_rows(conn, "users", 5)}
    schema_info = [
        SchemaInfo(table_name="users", columns=columns, row_count=total_rows)
    ]

    observation = SQLObservation(
        task_id="null_audit",
        task_description="Identify column in 'users' table with a null rate exceeding 50%",
        schema_info=schema_info,
        sample_rows=sample_rows,
        attempt_number=1,
        max_attempts=1,
        previous_feedback=None,
    )

    ground_truth = {
        "null_column": null_column,  # existing — the column with NULLs
        "all_columns": all_columns,  # NEW — for partial credit
        "table": "users",  # NEW — for D1/D2 query matching
        "column": null_column,  # NEW — alias for unified access
        "task_type": "null_audit",  # NEW — for compute_episode_score routing
    }

    return conn, observation, ground_truth


def create_type_repair_task(
    seed: int = 42,
) -> Tuple[sqlite3.Connection, SQLObservation, Dict[str, Any]]:
    """Create type repair task with transactions table."""
    import random

    rng = random.Random(seed)
    conn = sqlite3.connect(":memory:", check_same_thread=False)
    cursor = conn.cursor()

    cursor.execute(
        """
        CREATE TABLE transactions (
            id INTEGER PRIMARY KEY,
            user_id INTEGER NOT NULL,
            amount TEXT,
            description TEXT,
            category TEXT,
            date TEXT,
            transaction_date TEXT
        )
    """
    )

    cursor.execute(
        """
        CREATE TABLE categories (
            id INTEGER PRIMARY KEY,
            name TEXT NOT NULL,
            budget_limit REAL NOT NULL,
            created_at TEXT NOT NULL
        )
    """
    )

    for i, (name, budget) in enumerate(
        [
            ("Food & Beverage", 5000.0),
            ("Office Supplies", 2000.0),
            ("Travel", 8000.0),
            ("Software", 15000.0),
            ("Hardware", 25000.0),
            ("Marketing", 10000.0),
            ("HR", 6000.0),
            ("Legal", 12000.0),
        ],
        start=1,
    ):
        cursor.execute(
            "INSERT INTO categories (id, name, budget_limit, created_at) VALUES (?, ?, ?, ?)",
            (i, name, budget, "2024-01-01"),
        )

    transactions_data = []
    for i in range(1000):
        user_id = rng.randint(1, 100)
        amount_value = round(rng.uniform(10.0, 1000.0), 2)
        amount_text = f"{amount_value:.2f}"
        description = f"Transaction {i}"
        category = rng.choice(
            [
                "food",
                "software",
                "travel",
                "hardware",
                "marketing",
            ]
        )
        date = f"2023-{rng.randint(1, 12):02d}-{rng.randint(1, 28):02d}"
        transaction_date = f"2023-{rng.randint(1, 12):02d}-{rng.randint(1, 28):02d}"
        transactions_data.append(
            (i, user_id, amount_text, description, category, date, transaction_date)
        )

    cursor.executemany(
        """
        INSERT INTO transactions (
            id, user_id, amount, description, category, date, transaction_date
        )
        VALUES (?, ?, ?, ?, ?, ?, ?)
    """,
        transactions_data,
    )

    # Verify tables have data
    cursor.execute("SELECT COUNT(*) FROM transactions")
    count = cursor.fetchone()[0]
    assert count > 0, f"transactions table must have >0 rows, got {count}"
    
    cursor.execute("SELECT COUNT(*) FROM categories")
    count = cursor.fetchone()[0]
    assert count > 0, f"categories table must have >0 rows, got {count}"

    schema_info = []
    for table_name, row_count in [("transactions", 1000), ("categories", 8)]:
        cursor.execute(f"PRAGMA table_info({table_name})")
        columns = [f"{row[1]} ({row[2]})" for row in cursor.fetchall()]
        schema_info.append(
            SchemaInfo(table_name=table_name, columns=columns, row_count=row_count)
        )

    sample_rows = {
        "transactions": _extract_sample_rows(conn, "transactions", 5),
        "categories": _extract_sample_rows(conn, "categories", 5),
    }

    observation = SQLObservation(
        task_id="type_repair",
        task_description="A production data pipeline is reporting incorrect aggregate totals on financial records. Inspect all provided tables and sample rows to identify the data quality issue, then submit a fix_query that corrects it. Your fix must preserve all existing rows and their original values.",
        schema_info=schema_info,
        sample_rows=sample_rows,
        attempt_number=1,
        max_attempts=3,
        previous_feedback=None,
    )

    ground_truth = {
        "target_table": "transactions",  # existing
        "target_column": "amount",  # existing
        "current_type": "TEXT",  # existing
        "correct_type": "REAL",  # existing
        "sample_values": ["12.50", "345.75", "89.99"],  # existing
        "preserve_columns": ["id", "description", "amount"],  # compatible core fields
        "table": "transactions",  # NEW — unified alias
        "column": "amount",  # NEW — unified alias
        "task_type": "type_repair",  # NEW
    }

    return conn, observation, ground_truth


def create_dedup_bug_task(
    seed: int = 42,
) -> Tuple[sqlite3.Connection, SQLObservation, Dict[str, Any]]:
    """Create deduplication bug task with orders, order_items, and shipments."""
    import random

    rng = random.Random(seed)
    conn = sqlite3.connect(":memory:", check_same_thread=False)
    cursor = conn.cursor()

    cursor.execute(
        """
        CREATE TABLE orders (
            id INTEGER PRIMARY KEY,
            customer_id INTEGER NOT NULL,
            total_amount REAL NOT NULL
        )
    """
    )

    cursor.execute(
        """
        CREATE TABLE order_items (
            id INTEGER PRIMARY KEY,
            order_id INTEGER NOT NULL,
            product_name TEXT NOT NULL,
            quantity INTEGER NOT NULL,
            price REAL NOT NULL
        )
    """
    )

    cursor.execute(
        """
        CREATE TABLE shipments (
            id INTEGER PRIMARY KEY,
            order_item_id INTEGER NOT NULL,
            tracking_number TEXT NOT NULL,
            shipped_at TEXT NOT NULL
        )
    """
    )

    orders_data = []
    for i in range(200):
        customer_id = rng.randint(1, 50)
        total_amount = round(rng.uniform(50.0, 500.0), 2)
        orders_data.append((i, customer_id, total_amount))

    cursor.executemany(
        """
        INSERT INTO orders (id, customer_id, total_amount)
        VALUES (?, ?, ?)
    """,
        orders_data,
    )

    order_items_data = []
    for order_id in range(160):
        item_id = order_id
        product_name = f"Product {item_id}"
        quantity = rng.randint(1, 5)
        price = round(rng.uniform(10.0, 100.0), 2)
        order_items_data.append((item_id, order_id, product_name, quantity, price))

    cursor.executemany(
        """
        INSERT INTO order_items (id, order_id, product_name, quantity, price)
        VALUES (?, ?, ?, ?, ?)
    """,
        order_items_data,
    )

    duplicate_data = []
    for order_id in range(160):
        item_id = 160 + order_id
        product_name = f"Duplicate Product {item_id}"
        quantity = rng.randint(1, 5)
        price = round(rng.uniform(10.0, 100.0), 2)
        duplicate_data.append((item_id, order_id, product_name, quantity, price))

    cursor.executemany(
        """
        INSERT INTO order_items (id, order_id, product_name, quantity, price)
        VALUES (?, ?, ?, ?, ?)
    """,
        duplicate_data,
    )

    shipments_data = []
    for shipment_id in range(160):
        tracking_number = f"TRK{shipment_id:06d}"
        shipped_at = (
            f"2024-{rng.randint(1, 12):02d}-{rng.randint(1, 28):02d}"
            f"T{rng.randint(8, 20):02d}:{rng.randint(0, 59):02d}:00"
        )
        shipments_data.append((shipment_id, shipment_id, tracking_number, shipped_at))

    cursor.executemany(
        """
        INSERT INTO shipments (id, order_item_id, tracking_number, shipped_at)
        VALUES (?, ?, ?, ?)
    """,
        shipments_data,
    )

    # Verify tables have data
    cursor.execute("SELECT COUNT(*) FROM orders")
    count = cursor.fetchone()[0]
    assert count > 0, f"orders table must have >0 rows, got {count}"
    
    cursor.execute("SELECT COUNT(*) FROM order_items")
    count = cursor.fetchone()[0]
    assert count > 0, f"order_items table must have >0 rows, got {count}"
    
    cursor.execute("SELECT COUNT(*) FROM shipments")
    count = cursor.fetchone()[0]
    assert count > 0, f"shipments table must have >0 rows, got {count}"

    # Count legitimate rows (unique order_ids) for preservation grading
    clean_count_cursor = conn.execute(
        "SELECT COUNT(DISTINCT order_id) FROM order_items"
    )
    expected_clean_count = clean_count_cursor.fetchone()[0]
    original_row_count = conn.execute("SELECT COUNT(*) FROM order_items").fetchone()[0]

    schema_info = []
    for table_name, row_count in [
        ("orders", 200),
        ("order_items", 320),
        ("shipments", 160),
    ]:
        validated_name = _validate_table_name(table_name)
        cursor.execute(f"PRAGMA table_info({validated_name})")
        columns = [f"{row[1]} ({row[2]})" for row in cursor.fetchall()]
        schema_info.append(
            SchemaInfo(table_name=table_name, columns=columns, row_count=row_count)
        )

    sample_rows = {}
    sample_rows["orders"] = _extract_sample_rows(conn, "orders", 5)

    # Deterministic sampling: ensure at least one duplicate pair is visible
    # First, get one order_id that has duplicates
    cursor.execute("""
        SELECT order_id FROM order_items 
        GROUP BY order_id HAVING COUNT(*) > 1 
        ORDER BY order_id LIMIT 1
    """)
    dup_order_id = cursor.fetchone()
    
    if dup_order_id:
        dup_order_id = dup_order_id[0]
        # Get 2 rows with the duplicate order_id (to show the pattern)
        cursor.execute("SELECT * FROM order_items WHERE order_id = ? LIMIT 2", (dup_order_id,))
        dup_rows = cursor.fetchall()
        dup_columns = [description[0] for description in cursor.description]
        
        # Fill remaining samples with other rows (ensuring variety)
        remaining = 5 - len(dup_rows)
        cursor.execute("""
            SELECT * FROM order_items 
            WHERE order_id != ? 
            ORDER BY rowid LIMIT ?
        """, (dup_order_id, remaining))
        other_rows = cursor.fetchall()
        
        all_rows = dup_rows + other_rows
        sample_rows["order_items"] = [dict(zip(dup_columns, row)) for row in all_rows]
    else:
        # Fallback to deterministic order if no duplicates found (shouldn't happen)
        cursor.execute("SELECT * FROM order_items ORDER BY rowid LIMIT 5")
        columns = [description[0] for description in cursor.description]
        rows = cursor.fetchall()
        sample_rows["order_items"] = [dict(zip(columns, row)) for row in rows]
    
    sample_rows["shipments"] = _extract_sample_rows(conn, "shipments", 5)

    observation = SQLObservation(
        task_id="dedup_bug",
        task_description="A reporting query that joins two tables is returning more rows than expected based on known business volumes. Inspect the schema and sample data, identify which table contains the data integrity issue and which column is responsible, then submit a fix_query that resolves the duplication.",
        schema_info=schema_info,
        sample_rows=sample_rows,
        attempt_number=1,
        max_attempts=5,
        previous_feedback=None,
    )

    ground_truth = {
        "root_cause_table": "order_items",  # existing
        "offending_key": "order_id",  # existing
        "expected_clean_count": expected_clean_count,    # rows that SHOULD remain
        "original_row_count": original_row_count,      # rows before fix
        "duplicate_count": original_row_count - expected_clean_count,  # rows to delete
        "table": "order_items",  # NEW — unified alias
        "column": "order_id",  # NEW — unified alias
        "task_type": "dedup_bug",  # NEW
    }

    return conn, observation, ground_truth


def create_fk_integrity_task(
    seed: int = 42,
) -> Tuple[sqlite3.Connection, SQLObservation, Dict[str, Any]]:
    """Create FK integrity task with orphan child rows in invoices.customer_id."""
    import random

    rng = random.Random(seed)
    conn = sqlite3.connect(":memory:", check_same_thread=False)
    cursor = conn.cursor()

    cursor.execute(
        """
        CREATE TABLE customers (
            id INTEGER PRIMARY KEY,
            name TEXT NOT NULL,
            segment TEXT NOT NULL,
            created_at TEXT NOT NULL
        )
    """
    )
    cursor.execute(
        """
        CREATE TABLE invoices (
            id INTEGER PRIMARY KEY,
            customer_id INTEGER NOT NULL,
            amount REAL NOT NULL,
            status TEXT NOT NULL,
            issued_at TEXT NOT NULL
        )
    """
    )

    customer_rows = []
    for i in range(1, 401):
        customer_rows.append(
            (
                i,
                f"Customer_{i}",
                rng.choice(["enterprise", "mid_market", "smb"]),
                f"2024-{rng.randint(1,12):02d}-{rng.randint(1,28):02d}",
            )
        )
    cursor.executemany(
        "INSERT INTO customers (id, name, segment, created_at) VALUES (?, ?, ?, ?)",
        customer_rows,
    )

    invoice_rows = []
    for i in range(1, 1201):
        invoice_rows.append(
            (
                i,
                rng.randint(1, 400),
                round(rng.uniform(25.0, 2500.0), 2),
                rng.choice(["issued", "paid", "overdue"]),
                f"2025-{rng.randint(1,12):02d}-{rng.randint(1,28):02d}",
            )
        )

    orphan_invoice_ids = set(rng.sample(range(1, 1201), 180))
    rewritten_rows = []
    valid_invoice_ids = []
    for row in invoice_rows:
        invoice_id, customer_id, amount, status, issued_at = row
        if invoice_id in orphan_invoice_ids:
            rewritten_rows.append(
                (invoice_id, 10000 + rng.randint(1, 1000), amount, status, issued_at)
            )
        else:
            rewritten_rows.append(row)
            valid_invoice_ids.append(invoice_id)

    cursor.executemany(
        "INSERT INTO invoices (id, customer_id, amount, status, issued_at) VALUES (?, ?, ?, ?, ?)",
        rewritten_rows,
    )

    # Verify tables have data
    cursor.execute("SELECT COUNT(*) FROM customers")
    count = cursor.fetchone()[0]
    assert count > 0, f"customers table must have >0 rows, got {count}"
    
    cursor.execute("SELECT COUNT(*) FROM invoices")
    count = cursor.fetchone()[0]
    assert count > 0, f"invoices table must have >0 rows, got {count}"

    schema_info = []
    for table_name, row_count in [("customers", 400), ("invoices", 1200)]:
        cursor.execute(f"PRAGMA table_info({table_name})")
        columns = [f"{row[1]} ({row[2]})" for row in cursor.fetchall()]
        schema_info.append(
            SchemaInfo(table_name=table_name, columns=columns, row_count=row_count)
        )

    sample_rows = {
        "customers": _extract_sample_rows(conn, "customers", 5),
        "invoices": _extract_sample_rows(conn, "invoices", 5),
    }

    observation = SQLObservation(
        task_id="fk_integrity",
        task_description=(
            "A billing pipeline shows referential integrity drift between parent and child records. "
            "Inspect the schema and samples, then submit a fix_query that removes orphan relationships "
            "while preserving valid invoice records."
        ),
        schema_info=schema_info,
        sample_rows=sample_rows,
        attempt_number=1,
        max_attempts=4,
        previous_feedback=None,
    )

    ground_truth = {
        "task_type": "fk_integrity",
        "table": "invoices",
        "column": "customer_id",
        "parent_table": "customers",
        "parent_key": "id",
        "expected_clean_count": len(valid_invoice_ids),
        "valid_invoice_ids": valid_invoice_ids,
        "orphan_count": len(orphan_invoice_ids),
        "preserve_columns": ["id", "customer_id", "amount", "status", "issued_at"],
    }
    return conn, observation, ground_truth


def create_schema_drift_task(
    seed: int = 42,
) -> Tuple[sqlite3.Connection, SQLObservation, Dict[str, Any]]:
    """Create schema drift task with two upstream versions that need canonicalization."""
    import random

    rng = random.Random(seed)
    conn = sqlite3.connect(":memory:", check_same_thread=False)
    cursor = conn.cursor()

    cursor.execute(
        """
        CREATE TABLE events_v1 (
            id INTEGER PRIMARY KEY,
            user_id INTEGER NOT NULL,
            event_name TEXT NOT NULL,
            amount_text TEXT NOT NULL,
            event_time TEXT NOT NULL
        )
    """
    )
    cursor.execute(
        """
        CREATE TABLE events_v2 (
            event_id INTEGER PRIMARY KEY,
            account_id INTEGER NOT NULL,
            event_type TEXT NOT NULL,
            amount REAL NOT NULL,
            event_ts TEXT NOT NULL,
            source TEXT NOT NULL
        )
    """
    )

    v1_rows = []
    for i in range(1, 601):
        amount = round(rng.uniform(5.0, 500.0), 2)
        v1_rows.append(
            (
                i,
                rng.randint(1, 300),
                rng.choice(["purchase", "refund", "chargeback"]),
                f"{amount:.2f}",
                f"2025-{rng.randint(1,12):02d}-{rng.randint(1,28):02d}T{rng.randint(0,23):02d}:{rng.randint(0,59):02d}:00",
            )
        )
    cursor.executemany(
        "INSERT INTO events_v1 (id, user_id, event_name, amount_text, event_time) VALUES (?, ?, ?, ?, ?)",
        v1_rows,
    )

    v2_rows = []
    for i in range(601, 1201):
        v2_rows.append(
            (
                i,
                rng.randint(1, 300),
                rng.choice(["purchase", "refund", "chargeback"]),
                round(rng.uniform(5.0, 500.0), 2),
                f"2025-{rng.randint(1,12):02d}-{rng.randint(1,28):02d}T{rng.randint(0,23):02d}:{rng.randint(0,59):02d}:00",
                rng.choice(["batch_a", "batch_b"]),
            )
        )
    cursor.executemany(
        "INSERT INTO events_v2 (event_id, account_id, event_type, amount, event_ts, source) VALUES (?, ?, ?, ?, ?, ?)",
        v2_rows,
    )

    # Verify tables have data
    cursor.execute("SELECT COUNT(*) FROM events_v1")
    count = cursor.fetchone()[0]
    assert count > 0, f"events_v1 table must have >0 rows, got {count}"
    
    cursor.execute("SELECT COUNT(*) FROM events_v2")
    count = cursor.fetchone()[0]
    assert count > 0, f"events_v2 table must have >0 rows, got {count}"

    schema_info = []
    for table_name, row_count in [("events_v1", 600), ("events_v2", 600)]:
        cursor.execute(f"PRAGMA table_info({table_name})")
        columns = [f"{row[1]} ({row[2]})" for row in cursor.fetchall()]
        schema_info.append(
            SchemaInfo(table_name=table_name, columns=columns, row_count=row_count)
        )

    sample_rows = {
        "events_v1": _extract_sample_rows(conn, "events_v1", 5),
        "events_v2": _extract_sample_rows(conn, "events_v2", 5),
    }

    observation = SQLObservation(
        task_id="schema_drift",
        task_description=(
            "Two upstream event producers emit the same business concept with drifted schemas. "
            "Create a canonical_events artifact that normalizes both streams into one aligned schema "
            "with numeric amounts and consistent timestamps."
        ),
        schema_info=schema_info,
        sample_rows=sample_rows,
        attempt_number=1,
        max_attempts=4,
        previous_feedback=None,
    )

    ground_truth = {
        "task_type": "schema_drift",
        "table": "canonical_events",
        "column": "amount",
        "source_tables": ["events_v1", "events_v2"],
        "expected_columns": ["event_id", "user_id", "event_type", "amount", "event_time"],
        "expected_row_count": 1200,
        "preserve_columns": ["event_id", "user_id", "event_type", "amount", "event_time"],
    }
    return conn, observation, ground_truth