"""
tasks.py — The three SQL optimization tasks (easy → medium → hard).

Each task mirrors a real-world scenario a DBA or backend engineer would face:
  1. Easy   — SELECT * + missing WHERE index on an orders table
  2. Medium — Correlated subquery anti-pattern on a user analytics table
  3. Hard   — Inefficient aggregation with missing CTE / window function optimization
"""

from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any


@dataclass
class Task:
    task_id: str
    difficulty: str          # easy | medium | hard
    description: str
    schema_ddl: str
    slow_query: str
    reference_query: str     # canonical optimized solution (for grading)
    execution_stats: dict[str, Any]
    # Structural signals the grader checks for in the optimized query
    expected_improvements: list[str] = field(default_factory=list)
    forbidden_patterns: list[str] = field(default_factory=list)
    # Maximum estimated_cost for full performance score
    target_cost: float = 100.0


# ---------------------------------------------------------------------------
# Task 1 — Easy: SELECT * + no index on filtered column
# ---------------------------------------------------------------------------

TASK_EASY = Task(
    task_id="simple_select_optimization",
    difficulty="easy",
    description=(
        "Optimize a query that uses SELECT * and filters on an un-indexed column. "
        "The agent should replace SELECT * with specific columns and suggest an index hint."
    ),
    schema_ddl="""
-- E-commerce orders table (50M rows in production)
CREATE TABLE orders (
    order_id     BIGINT PRIMARY KEY,
    customer_id  BIGINT NOT NULL,
    status       VARCHAR(20) NOT NULL,   -- 'pending','shipped','delivered','cancelled'
    total_amount DECIMAL(12,2) NOT NULL,
    created_at   TIMESTAMP NOT NULL,
    updated_at   TIMESTAMP NOT NULL,
    shipping_address TEXT,
    notes        TEXT
);

-- Existing indexes
CREATE INDEX idx_orders_customer ON orders(customer_id);
-- NOTE: No index on (status, created_at)
""",
    slow_query="""
SELECT *
FROM orders
WHERE status = 'pending'
  AND created_at >= '2024-01-01'
ORDER BY created_at DESC;
""",
    reference_query="""
SELECT order_id, customer_id, status, total_amount, created_at
FROM orders
WHERE status = 'pending'
  AND created_at >= '2024-01-01'
ORDER BY created_at DESC;
-- Recommended: CREATE INDEX idx_orders_status_created ON orders(status, created_at DESC);
""",
    execution_stats={
        "estimated_cost": 9800.0,
        "scan_type": "seq_scan",
        "rows_examined": 50_000_000,
        "has_anti_patterns": ["SELECT *", "seq_scan on large table"],
    },
    expected_improvements=["specific columns instead of *", "index hint or note"],
    forbidden_patterns=["SELECT \\*"],
    target_cost=200.0,
)


# ---------------------------------------------------------------------------
# Task 2 — Medium: Correlated subquery anti-pattern
# ---------------------------------------------------------------------------

TASK_MEDIUM = Task(
    task_id="join_optimization",
    difficulty="medium",
    description=(
        "Rewrite a correlated subquery that runs once per row into an efficient "
        "JOIN or CTE. The query finds top customers who spent more than the average "
        "order value in their region."
    ),
    schema_ddl="""
CREATE TABLE customers (
    customer_id  BIGINT PRIMARY KEY,
    name         VARCHAR(255) NOT NULL,
    email        VARCHAR(255) UNIQUE NOT NULL,
    region       VARCHAR(50) NOT NULL,
    created_at   TIMESTAMP NOT NULL
);

CREATE TABLE orders (
    order_id     BIGINT PRIMARY KEY,
    customer_id  BIGINT NOT NULL REFERENCES customers(customer_id),
    region       VARCHAR(50) NOT NULL,
    total_amount DECIMAL(12,2) NOT NULL,
    status       VARCHAR(20) NOT NULL,
    created_at   TIMESTAMP NOT NULL
);

CREATE INDEX idx_orders_customer   ON orders(customer_id);
CREATE INDEX idx_orders_region     ON orders(region);
CREATE INDEX idx_customers_region  ON customers(region);
""",
    slow_query="""
SELECT
    c.customer_id,
    c.name,
    c.region,
    SUM(o.total_amount) AS total_spent
FROM customers c
JOIN orders o ON o.customer_id = c.customer_id
WHERE o.status = 'delivered'
  AND o.total_amount > (
      SELECT AVG(o2.total_amount)
      FROM orders o2
      WHERE o2.region = o.region
        AND o2.status = 'delivered'
  )
GROUP BY c.customer_id, c.name, c.region
ORDER BY total_spent DESC
LIMIT 100;
""",
    reference_query="""
WITH regional_avg AS (
    SELECT region, AVG(total_amount) AS avg_amount
    FROM orders
    WHERE status = 'delivered'
    GROUP BY region
)
SELECT
    c.customer_id,
    c.name,
    c.region,
    SUM(o.total_amount) AS total_spent
FROM customers c
JOIN orders o ON o.customer_id = c.customer_id
JOIN regional_avg ra ON ra.region = o.region
WHERE o.status = 'delivered'
  AND o.total_amount > ra.avg_amount
GROUP BY c.customer_id, c.name, c.region
ORDER BY total_spent DESC
LIMIT 100;
""",
    execution_stats={
        "estimated_cost": 48_750.0,
        "scan_type": "seq_scan",
        "rows_examined": 12_000_000,
        "has_anti_patterns": [
            "correlated subquery (runs once per row)",
            "missing CTE for repeated aggregation",
        ],
    },
    expected_improvements=["CTE or derived table for regional_avg", "JOIN instead of correlated subquery"],
    forbidden_patterns=["SELECT.*AVG.*FROM.*WHERE.*=.*\\.region"],  # correlated subquery pattern
    target_cost=1_500.0,
)


# ---------------------------------------------------------------------------
# Task 3 — Hard: Complex aggregation with window functions needed
# ---------------------------------------------------------------------------

TASK_HARD = Task(
    task_id="aggregation_optimization",
    difficulty="hard",
    description=(
        "Optimize a complex monthly cohort revenue report that uses repeated "
        "self-joins and subqueries where window functions and a single-pass CTE "
        "pipeline would be dramatically faster."
    ),
    schema_ddl="""
CREATE TABLE events (
    event_id    BIGINT PRIMARY KEY,
    user_id     BIGINT NOT NULL,
    event_type  VARCHAR(50) NOT NULL,   -- 'purchase','refund','view','signup'
    amount      DECIMAL(12,2),
    occurred_at TIMESTAMP NOT NULL,
    session_id  UUID,
    metadata    JSONB
);

CREATE TABLE users (
    user_id      BIGINT PRIMARY KEY,
    signup_date  DATE NOT NULL,
    plan         VARCHAR(20) NOT NULL,  -- 'free','pro','enterprise'
    country      VARCHAR(3) NOT NULL
);

CREATE INDEX idx_events_user       ON events(user_id);
CREATE INDEX idx_events_type_date  ON events(event_type, occurred_at);
CREATE INDEX idx_users_signup      ON users(signup_date);
CREATE INDEX idx_users_plan        ON users(plan);
""",
    slow_query="""
SELECT
    DATE_TRUNC('month', u.signup_date) AS cohort_month,
    u.plan,
    COUNT(DISTINCT u.user_id) AS cohort_size,
    (
        SELECT COALESCE(SUM(e.amount), 0)
        FROM events e
        WHERE e.user_id = u.user_id
          AND e.event_type = 'purchase'
          AND e.occurred_at >= DATE_TRUNC('month', u.signup_date)
          AND e.occurred_at <  DATE_TRUNC('month', u.signup_date) + INTERVAL '30 days'
    ) AS m0_revenue,
    (
        SELECT COALESCE(SUM(e.amount), 0)
        FROM events e
        WHERE e.user_id = u.user_id
          AND e.event_type = 'purchase'
          AND e.occurred_at >= DATE_TRUNC('month', u.signup_date) + INTERVAL '30 days'
          AND e.occurred_at <  DATE_TRUNC('month', u.signup_date) + INTERVAL '60 days'
    ) AS m1_revenue,
    (
        SELECT COALESCE(SUM(e.amount), 0)
        FROM events e
        WHERE e.user_id = u.user_id
          AND e.event_type = 'purchase'
          AND e.occurred_at >= DATE_TRUNC('month', u.signup_date) + INTERVAL '60 days'
          AND e.occurred_at <  DATE_TRUNC('month', u.signup_date) + INTERVAL '90 days'
    ) AS m2_revenue
FROM users u
WHERE u.signup_date >= '2023-01-01'
GROUP BY cohort_month, u.plan, u.user_id, u.signup_date
ORDER BY cohort_month, u.plan;
""",
    reference_query="""
WITH user_purchases AS (
    SELECT
        e.user_id,
        e.amount,
        e.occurred_at,
        u.signup_date,
        u.plan,
        DATE_TRUNC('month', u.signup_date) AS cohort_month,
        (e.occurred_at - u.signup_date::TIMESTAMP) AS days_since_signup
    FROM users u
    JOIN events e ON e.user_id = u.user_id
    WHERE u.signup_date >= '2023-01-01'
      AND e.event_type = 'purchase'
),
cohort_revenue AS (
    SELECT
        cohort_month,
        plan,
        user_id,
        SUM(amount) FILTER (WHERE days_since_signup < INTERVAL '30 days')  AS m0_revenue,
        SUM(amount) FILTER (WHERE days_since_signup >= INTERVAL '30 days'
                              AND days_since_signup <  INTERVAL '60 days') AS m1_revenue,
        SUM(amount) FILTER (WHERE days_since_signup >= INTERVAL '60 days'
                              AND days_since_signup <  INTERVAL '90 days') AS m2_revenue
    FROM user_purchases
    GROUP BY cohort_month, plan, user_id
)
SELECT
    cohort_month,
    plan,
    COUNT(DISTINCT user_id)  AS cohort_size,
    COALESCE(SUM(m0_revenue), 0) AS m0_revenue,
    COALESCE(SUM(m1_revenue), 0) AS m1_revenue,
    COALESCE(SUM(m2_revenue), 0) AS m2_revenue
FROM cohort_revenue
GROUP BY cohort_month, plan
ORDER BY cohort_month, plan;
""",
    execution_stats={
        "estimated_cost": 892_000.0,
        "scan_type": "seq_scan",
        "rows_examined": 450_000_000,
        "has_anti_patterns": [
            "3x correlated scalar subqueries",
            "GROUP BY includes user_id causing over-aggregation",
            "no FILTER aggregate — using WHERE in subqueries",
        ],
    },
    expected_improvements=[
        "CTE to pre-join events once",
        "FILTER aggregate or CASE WHEN instead of correlated subqueries",
        "single GROUP BY at cohort level",
    ],
    forbidden_patterns=[],
    target_cost=15_000.0,
)


TASKS: dict[str, Task] = {
    TASK_EASY.task_id:   TASK_EASY,
    TASK_MEDIUM.task_id: TASK_MEDIUM,
    TASK_HARD.task_id:   TASK_HARD,
}