"""
Task definitions for the SQL Debug Environment.
Each task has:
- schema_sql: CREATE TABLE statements
- sample_data: INSERT statements
- buggy_query: broken SQL
- error_message: error when running buggy_query (or "" for semantic bugs)
- task_description: what the query should do
- expected_columns: columns the correct result must have
- correct_query: reference solution (used only by grader)
- expected_rows: sorted list of tuples the correct query returns
- difficulty: easy | medium | hard
"""

from typing import Optional

TASKS = {
    "easy_missing_from": {
        "difficulty": "easy",
        "task_description": (
            "List the first name and last name of all employees in department 10, "
            "ordered by last name ascending."
        ),
        "schema_sql": """
CREATE TABLE employees (
    id INTEGER PRIMARY KEY,
    first_name TEXT NOT NULL,
    last_name TEXT NOT NULL,
    department_id INTEGER NOT NULL,
    salary REAL NOT NULL
);
""",
        "sample_data": """
INSERT INTO employees VALUES (1, 'Alice', 'Anderson', 10, 55000);
INSERT INTO employees VALUES (2, 'Bob', 'Brown', 20, 62000);
INSERT INTO employees VALUES (3, 'Carol', 'Clark', 10, 71000);
INSERT INTO employees VALUES (4, 'David', 'Davis', 10, 48000);
INSERT INTO employees VALUES (5, 'Eve', 'Evans', 30, 95000);
""",
        "buggy_query": "SELECT first_name, last_name employees WHERE department_id = 10 ORDER BY last_name;",
        "error_message": 'SyntaxError: near "employees": syntax error',
        "expected_columns": ["first_name", "last_name"],
        "correct_query": "SELECT first_name, last_name FROM employees WHERE department_id = 10 ORDER BY last_name;",
        "expected_rows": [
            ("Alice", "Anderson"),
            ("Carol", "Clark"),
            ("David", "Davis"),
        ],
    },
    "easy_wrong_column": {
        "difficulty": "easy",
        "task_description": (
            "Find the total salary budget for each department. "
            "Return department_id and the total salary, ordered by department_id."
        ),
        "schema_sql": """
CREATE TABLE employees (
    id INTEGER PRIMARY KEY,
    first_name TEXT NOT NULL,
    last_name TEXT NOT NULL,
    department_id INTEGER NOT NULL,
    salary REAL NOT NULL
);
""",
        "sample_data": """
INSERT INTO employees VALUES (1, 'Alice', 'Anderson', 10, 55000);
INSERT INTO employees VALUES (2, 'Bob', 'Brown', 20, 62000);
INSERT INTO employees VALUES (3, 'Carol', 'Clark', 10, 71000);
INSERT INTO employees VALUES (4, 'David', 'Davis', 10, 48000);
INSERT INTO employees VALUES (5, 'Eve', 'Evans', 30, 95000);
""",
        "buggy_query": "SELECT department_id, SUM(wage) AS total_salary FROM employees GROUP BY department_id ORDER BY department_id;",
        "error_message": "OperationalError: no such column: wage",
        "expected_columns": ["department_id", "total_salary"],
        "correct_query": "SELECT department_id, SUM(salary) AS total_salary FROM employees GROUP BY department_id ORDER BY department_id;",
        "expected_rows": [(10, 174000.0), (20, 62000.0), (30, 95000.0)],
    },
    "easy_missing_groupby": {
        "difficulty": "easy",
        "task_description": (
            "Count the number of orders placed by each customer. "
            "Return customer_id and order_count, only for customers with more than 1 order."
        ),
        "schema_sql": """
CREATE TABLE orders (
    id INTEGER PRIMARY KEY,
    customer_id INTEGER NOT NULL,
    product TEXT NOT NULL,
    amount REAL NOT NULL,
    order_date TEXT NOT NULL
);
""",
        "sample_data": """
INSERT INTO orders VALUES (1, 101, 'Laptop', 999.99, '2024-01-15');
INSERT INTO orders VALUES (2, 102, 'Mouse', 29.99, '2024-01-16');
INSERT INTO orders VALUES (3, 101, 'Keyboard', 79.99, '2024-01-17');
INSERT INTO orders VALUES (4, 103, 'Monitor', 349.99, '2024-01-18');
INSERT INTO orders VALUES (5, 101, 'Webcam', 89.99, '2024-01-19');
INSERT INTO orders VALUES (6, 102, 'Headset', 149.99, '2024-01-20');
""",
        "buggy_query": "SELECT customer_id, COUNT(*) AS order_count FROM orders HAVING COUNT(*) > 1;",
        "error_message": "OperationalError: HAVING clause on a non-aggregate query",
        "expected_columns": ["customer_id", "order_count"],
        "correct_query": "SELECT customer_id, COUNT(*) AS order_count FROM orders GROUP BY customer_id HAVING COUNT(*) > 1 ORDER BY customer_id;",
        "expected_rows": [(101, 3), (102, 2)],
    },
    "medium_wrong_join": {
        "difficulty": "medium",
        "task_description": (
            "List each employee's full name and their department name. "
            "Return first_name, last_name, department_name ordered by last_name."
        ),
        "schema_sql": """
CREATE TABLE employees (
    id INTEGER PRIMARY KEY,
    first_name TEXT NOT NULL,
    last_name TEXT NOT NULL,
    dept_id INTEGER NOT NULL,
    salary REAL NOT NULL
);

CREATE TABLE departments (
    id INTEGER PRIMARY KEY,
    name TEXT NOT NULL,
    budget REAL NOT NULL
);
""",
        "sample_data": """
INSERT INTO departments VALUES (1, 'Engineering', 500000);
INSERT INTO departments VALUES (2, 'Marketing', 200000);
INSERT INTO departments VALUES (3, 'HR', 150000);

INSERT INTO employees VALUES (1, 'Alice', 'Anderson', 1, 90000);
INSERT INTO employees VALUES (2, 'Bob', 'Brown', 2, 70000);
INSERT INTO employees VALUES (3, 'Carol', 'Clark', 1, 85000);
INSERT INTO employees VALUES (4, 'Diana', 'Dole', 3, 65000);
""",
        "buggy_query": """
SELECT e.first_name, e.last_name, d.name AS department_name
FROM employees e
JOIN departments d ON e.id = d.id
ORDER BY e.last_name;
""",
        "error_message": "",
        "expected_columns": ["first_name", "last_name", "department_name"],
        "correct_query": """
SELECT e.first_name, e.last_name, d.name AS department_name
FROM employees e
JOIN departments d ON e.dept_id = d.id
ORDER BY e.last_name;
""",
        "expected_rows": [
            ("Alice", "Anderson", "Engineering"),
            ("Bob", "Brown", "Marketing"),
            ("Carol", "Clark", "Engineering"),
            ("Diana", "Dole", "HR"),
        ],
    },
    "medium_wrong_aggregate": {
        "difficulty": "medium",
        "task_description": (
            "Find the average order amount per customer. "
            "Return customer_id and avg_amount for customers whose average is above 100. "
            "Order by avg_amount descending."
        ),
        "schema_sql": """
CREATE TABLE orders (
    id INTEGER PRIMARY KEY,
    customer_id INTEGER NOT NULL,
    product TEXT NOT NULL,
    amount REAL NOT NULL
);
""",
        "sample_data": """
INSERT INTO orders VALUES (1, 101, 'Laptop', 999.99);
INSERT INTO orders VALUES (2, 102, 'Mouse', 29.99);
INSERT INTO orders VALUES (3, 101, 'Keyboard', 79.99);
INSERT INTO orders VALUES (4, 103, 'Monitor', 349.99);
INSERT INTO orders VALUES (5, 101, 'Webcam', 89.99);
INSERT INTO orders VALUES (6, 102, 'Headset', 149.99);
INSERT INTO orders VALUES (7, 103, 'Desk', 599.99);
""",
        "buggy_query": """
SELECT customer_id, SUM(amount) AS avg_amount
FROM orders
GROUP BY customer_id
HAVING avg_amount > 100
ORDER BY avg_amount DESC;
""",
        "error_message": "",
        "expected_columns": ["customer_id", "avg_amount"],
        "correct_query": """
SELECT customer_id, AVG(amount) AS avg_amount
FROM orders
GROUP BY customer_id
HAVING AVG(amount) > 100
ORDER BY avg_amount DESC;
""",
        "expected_rows": [(103, 474.99), (101, 389.99)],
    },
    "medium_inner_vs_left": {
        "difficulty": "medium",
        "task_description": (
            "List ALL customers along with the total amount they've spent. "
            "Customers who haven't placed any orders should show 0.0 as total_spent. "
            "Return customer_id, customer_name, total_spent ordered by customer_id."
        ),
        "schema_sql": """
CREATE TABLE customers (
    id INTEGER PRIMARY KEY,
    name TEXT NOT NULL,
    email TEXT NOT NULL
);

CREATE TABLE orders (
    id INTEGER PRIMARY KEY,
    customer_id INTEGER NOT NULL,
    amount REAL NOT NULL
);
""",
        "sample_data": """
INSERT INTO customers VALUES (1, 'Alice', 'alice@example.com');
INSERT INTO customers VALUES (2, 'Bob', 'bob@example.com');
INSERT INTO customers VALUES (3, 'Carol', 'carol@example.com');
INSERT INTO customers VALUES (4, 'Dave', 'dave@example.com');

INSERT INTO orders VALUES (1, 1, 250.00);
INSERT INTO orders VALUES (2, 1, 150.00);
INSERT INTO orders VALUES (3, 3, 400.00);
""",
        "buggy_query": """
SELECT c.id AS customer_id, c.name AS customer_name, SUM(o.amount) AS total_spent
FROM customers c
JOIN orders o ON c.id = o.customer_id
GROUP BY c.id, c.name
ORDER BY c.id;
""",
        "error_message": "",
        "expected_columns": ["customer_id", "customer_name", "total_spent"],
        "correct_query": """
SELECT c.id AS customer_id, c.name AS customer_name, COALESCE(SUM(o.amount), 0.0) AS total_spent
FROM customers c
LEFT JOIN orders o ON c.id = o.customer_id
GROUP BY c.id, c.name
ORDER BY c.id;
""",
        "expected_rows": [
            (1, "Alice", 400.0),
            (2, "Bob", 0.0),
            (3, "Carol", 400.0),
            (4, "Dave", 0.0),
        ],
    },
    "hard_window_function": {
        "difficulty": "hard",
        "task_description": (
            "For each employee, compute a running_avg_salary: the average of the current "
            "employee's salary and the employee immediately before them in the same department "
            "(ordered by id). For the first employee in each department, it equals their own salary. "
            "Return dept_id, first_name, id, salary, running_avg_salary ordered by dept_id, id."
        ),
        "schema_sql": """
CREATE TABLE employees (
    id INTEGER PRIMARY KEY,
    first_name TEXT NOT NULL,
    dept_id INTEGER NOT NULL,
    salary REAL NOT NULL
);
""",
        "sample_data": """
INSERT INTO employees VALUES (1, 'Alice', 1, 80000);
INSERT INTO employees VALUES (2, 'Bob', 1, 90000);
INSERT INTO employees VALUES (3, 'Carol', 2, 70000);
INSERT INTO employees VALUES (4, 'Diana', 2, 60000);
INSERT INTO employees VALUES (5, 'Eve', 2, 80000);
""",
        "buggy_query": """
WITH lagged AS (
    SELECT id, first_name, dept_id, salary,
           LAG(salary, 2) OVER (PARTITION BY dept_id ORDER BY id) AS prev_salary
    FROM employees
)
SELECT dept_id, first_name, id, salary,
       ROUND((salary + COALESCE(prev_salary, salary)) / 2.0, 2) AS running_avg_salary
FROM lagged
ORDER BY dept_id, id;
""",
        "error_message": "",
        "expected_columns": [
            "dept_id",
            "first_name",
            "id",
            "salary",
            "running_avg_salary",
        ],
        "correct_query": """
WITH lagged AS (
    SELECT id, first_name, dept_id, salary,
           LAG(salary, 1) OVER (PARTITION BY dept_id ORDER BY id) AS prev_salary
    FROM employees
)
SELECT dept_id, first_name, id, salary,
       ROUND((salary + COALESCE(prev_salary, salary)) / 2.0, 2) AS running_avg_salary
FROM lagged
ORDER BY dept_id, id;
""",
        "expected_rows": [
            (1, "Alice", 1, 80000.0, 80000.0),
            (1, "Bob", 2, 90000.0, 85000.0),
            (2, "Carol", 3, 70000.0, 70000.0),
            (2, "Diana", 4, 60000.0, 65000.0),
            (2, "Eve", 5, 80000.0, 70000.0),
        ],
    },
    "hard_correlated_subquery": {
        "difficulty": "hard",
        "task_description": (
            "Find all products that have never been ordered. "
            "Return product_id and product_name ordered by product_id."
        ),
        "schema_sql": """
CREATE TABLE products (
    id INTEGER PRIMARY KEY,
    name TEXT NOT NULL,
    price REAL NOT NULL,
    category TEXT NOT NULL
);

CREATE TABLE order_items (
    id INTEGER PRIMARY KEY,
    order_id INTEGER NOT NULL,
    product_id INTEGER NOT NULL,
    quantity INTEGER NOT NULL
);
""",
        "sample_data": """
INSERT INTO products VALUES (1, 'Laptop', 999.99, 'Electronics');
INSERT INTO products VALUES (2, 'Mouse', 29.99, 'Electronics');
INSERT INTO products VALUES (3, 'Desk', 299.99, 'Furniture');
INSERT INTO products VALUES (4, 'Chair', 199.99, 'Furniture');
INSERT INTO products VALUES (5, 'Keyboard', 79.99, 'Electronics');

INSERT INTO order_items VALUES (1, 101, 1, 2);
INSERT INTO order_items VALUES (2, 101, 3, 1);
INSERT INTO order_items VALUES (3, 102, 2, 4);
""",
        "buggy_query": """
SELECT p.id AS product_id, p.name AS product_name
FROM products p
WHERE EXISTS (
    SELECT 1 FROM order_items oi WHERE oi.product_id = p.id
)
ORDER BY p.id;
""",
        "error_message": "",
        "expected_columns": ["product_id", "product_name"],
        "correct_query": """
SELECT p.id AS product_id, p.name AS product_name
FROM products p
WHERE NOT EXISTS (
    SELECT 1 FROM order_items oi WHERE oi.product_id = p.id
)
ORDER BY p.id;
""",
        "expected_rows": [(4, "Chair"), (5, "Keyboard")],
    },
    "hard_multi_cte": {
        "difficulty": "hard",
        "task_description": (
            "Calculate month-over-month revenue change. "
            "For each month (in YYYY-MM format), return the month, total_revenue, "
            "and revenue_change (current month minus previous month). "
            "For the first month, revenue_change should be 0. "
            "Order by month ascending."
        ),
        "schema_sql": """
CREATE TABLE sales (
    id INTEGER PRIMARY KEY,
    sale_date TEXT NOT NULL,
    amount REAL NOT NULL,
    region TEXT NOT NULL
);
""",
        "sample_data": """
INSERT INTO sales VALUES (1, '2024-01-05', 10000, 'North');
INSERT INTO sales VALUES (2, '2024-01-20', 15000, 'South');
INSERT INTO sales VALUES (3, '2024-02-08', 20000, 'North');
INSERT INTO sales VALUES (4, '2024-02-14', 12000, 'West');
INSERT INTO sales VALUES (5, '2024-03-01', 18000, 'South');
INSERT INTO sales VALUES (6, '2024-03-22', 22000, 'North');
""",
        "buggy_query": """
WITH monthly AS (
    SELECT strftime('%Y-%m', sale_date) AS month,
           SUM(amount) AS total_revenue
    FROM sales
    GROUP BY month
),
with_prev AS (
    SELECT month, total_revenue,
           LAG(total_revenue, 2) OVER (ORDER BY month) AS prev_revenue
    FROM monthly
)
SELECT month, total_revenue,
       COALESCE(total_revenue - prev_revenue, 0) AS revenue_change
FROM with_prev
ORDER BY month;
""",
        "error_message": "",
        "expected_columns": ["month", "total_revenue", "revenue_change"],
        "correct_query": """
WITH monthly AS (
    SELECT strftime('%Y-%m', sale_date) AS month,
           SUM(amount) AS total_revenue
    FROM sales
    GROUP BY month
),
with_prev AS (
    SELECT month, total_revenue,
           LAG(total_revenue, 1) OVER (ORDER BY month) AS prev_revenue
    FROM monthly
)
SELECT month, total_revenue,
       COALESCE(total_revenue - prev_revenue, 0) AS revenue_change
FROM with_prev
ORDER BY month;
""",
        "expected_rows": [
            ("2024-01", 25000.0, 0.0),
            ("2024-02", 32000.0, 7000.0),
            ("2024-03", 40000.0, 8000.0),
        ],
    },
}


def get_tasks_by_difficulty(difficulty: str) -> dict:
    return {k: v for k, v in TASKS.items() if v["difficulty"] == difficulty}


def get_task(task_id: str) -> Optional[dict]:
    return TASKS.get(task_id)


TASK_GROUPS = {
    "easy_syntax_fix": [
        "easy_missing_from",
        "easy_wrong_column",
        "easy_missing_groupby",
    ],
    "medium_logic_fix": [
        "medium_wrong_join",
        "medium_wrong_aggregate",
        "medium_inner_vs_left",
    ],
    "hard_optimization_fix": [
        "hard_window_function",
        "hard_correlated_subquery",
        "hard_multi_cte",
    ],
}