"""
Solution: the safety and grading code around a text-to-SQL model (Lab 04, 06_text_to_sql.py).
"""
import re
import sqlite3

FORBIDDEN = re.compile(r"\b(insert|update|delete|drop|alter|create|replace|attach|pragma|vacuum)\b", re.I)
CHECKED = ("status", "category", "name", "country")


def is_safe(sql: str) -> tuple[bool, str]:
    if ";" in sql.strip().rstrip(";"):
        return False, "more than one statement"
    if not re.match(r"^\s*(select|with)\b", sql, re.I):
        return False, "not a SELECT query"
    m = FORBIDDEN.search(sql)
    if m:
        return False, f"forbidden keyword: {m.group(0)}"
    return True, "ok"


def normalize(value) -> str:
    return f"{value:.2f}" if isinstance(value, float) else str(value)


def same_result(gold: list[tuple], predicted: list[tuple]) -> bool:
    if len(gold) != len(predicted):
        return False
    remaining = [[normalize(v) for v in row] for row in predicted]
    for row in gold:
        wanted = [normalize(v) for v in row]
        match = next((r for r in remaining if all(w in r for w in wanted)), None)
        if match is None:
            return False
        remaining.remove(match)
    return True


def unknown_values(db: sqlite3.Connection, sql: str) -> list[str]:
    found = []
    for col, op in re.findall(r"(?:\w+\.)?(status|category|name|country)\s*(=\s*'[^']*'|IN\s*\([^)]*\))", sql, re.I):
        found += [(col.lower(), v) for v in re.findall(r"'([^']*)'", op)]
    known: dict[str, set] = {}
    tables = [t for (t,) in db.execute("SELECT name FROM sqlite_master WHERE type = 'table'")]
    for t in tables:
        for c in db.execute(f"PRAGMA table_info({t})"):
            if c[1] in CHECKED:
                known.setdefault(c[1], set()).update(v for (v,) in db.execute(f"SELECT DISTINCT {c[1]} FROM {t}"))
    return [v for col, v in found if v not in known.get(col, set())]


def run_read_only(path: str, sql: str) -> list[tuple]:
    db = sqlite3.connect(f"file:{path}?mode=ro", uri=True)
    try:
        return db.execute(sql).fetchall()
    finally:
        db.close()