"""
Deterministic graders for the SQL Debug Environment.
Each grader scores an agent's submitted query from 0.0 to 1.0.

Scoring rubric:
  1.0 — Query executes without error AND returns exactly the expected rows
  0.7 — Query executes without error AND returns correct columns but wrong data
  0.5 — Query executes without error AND has correct columns
  0.3 — Query executes without error but has wrong columns
  0.0 — Query fails with an error
"""

import sqlite3
from typing import Tuple, List
from tasks import get_task


def normalize_rows(rows: List[Tuple]) -> List[Tuple]:
    """Normalize rows for comparison: round floats, lowercase strings."""
    normalized = []
    for row in rows:
        new_row = []
        for cell in row:
            if isinstance(cell, float):
                new_row.append(round(cell, 2))
            elif isinstance(cell, str):
                new_row.append(cell.strip())
            else:
                new_row.append(cell)
        normalized.append(tuple(new_row))
    return sorted(normalized)


def run_query_on_task(
    task_id: str, query: str
) -> Tuple[bool, str, List[Tuple], List[str]]:
    """
    Run a query against the task's in-memory SQLite database.

    Returns:
        (success, error_message, rows, column_names)
    """
    task = get_task(task_id)
    if task is None:
        return False, f"Unknown task_id: {task_id}", [], []

    conn = sqlite3.connect(":memory:")
    cursor = conn.cursor()
    try:

        cursor.executescript(task["schema_sql"])

        cursor.executescript(task["sample_data"])
        conn.commit()

        cursor.execute(query)
        rows = cursor.fetchall()
        col_names = (
            [desc[0] for desc in cursor.description] if cursor.description else []
        )
        return True, "", rows, col_names
    except Exception as e:
        return False, str(e), [], []
    finally:
        conn.close()


def grade_submission(task_id: str, submitted_query: str) -> Tuple[float, str, dict]:
    """
    Grade a submitted query for the given task.

    Returns:
        (score: float 0.0-1.0, feedback: str, details: dict)
    """
    task = get_task(task_id)
    if task is None:
        return 0.0, f"Unknown task: {task_id}", {}

    success, error, rows, col_names = run_query_on_task(task_id, submitted_query)

    expected_rows = task["expected_rows"]
    expected_cols = task["expected_columns"]

    details = {
        "task_id": task_id,
        "success": success,
        "error": error,
        "returned_rows": len(rows),
        "expected_rows": len(expected_rows),
        "returned_cols": col_names,
        "expected_cols": expected_cols,
    }

    if not success:
        return 0.0, f"Query failed with error: {error}", details

    cols_correct = sorted(col_names) == sorted(expected_cols)

    if not cols_correct:
        return (
            0.3,
            (
                f"Query runs but wrong columns. "
                f"Got {col_names}, expected {expected_cols}."
            ),
            details,
        )

    actual_normalized = normalize_rows(rows)
    expected_normalized = normalize_rows(expected_rows)

    if actual_normalized == expected_normalized:
        return (
            1.0,
            "Perfect! Query is correct and returns the expected result.",
            details,
        )
    elif len(actual_normalized) == len(expected_normalized):
        return (
            0.7,
            (
                f"Correct number of rows ({len(rows)}) but data doesn't match. "
                f"First actual row: {rows[0] if rows else 'none'}"
            ),
            details,
        )
    else:
        return (
            0.5,
            (
                f"Query runs and has correct columns but wrong number of rows. "
                f"Got {len(rows)}, expected {len(expected_rows)}."
            ),
            details,
        )


def compute_episode_score(task_id: str, submissions: List[str]) -> float:
    """
    Compute a final episode score [0.0, 1.0] across all attempts.
    Awards partial credit for early correctness and penalizes wasted attempts.
    """
    if not submissions:
        return 0.0

    best_score = 0.0
    for i, query in enumerate(submissions):
        score, _, _ = grade_submission(task_id, query)

        attempt_bonus = max(0, (5 - i) * 0.05)
        adjusted = min(1.0, score + attempt_bonus if score == 1.0 else score)
        best_score = max(best_score, adjusted)

    return round(best_score, 4)


def grade_easy(task_subid: str, submitted_query: str) -> float:
    """Grade a single easy task submission. Returns 0.0-1.0."""
    score, _, _ = grade_submission(task_subid, submitted_query)
    return score


def grade_medium(task_subid: str, submitted_query: str) -> float:
    """Grade a single medium task submission. Returns 0.0-1.0."""
    score, _, _ = grade_submission(task_subid, submitted_query)
    return score


def grade_hard(task_subid: str, submitted_query: str) -> float:
    """Grade a single hard task submission. Returns 0.0-1.0."""
    score, _, _ = grade_submission(task_subid, submitted_query)
    return score