import sqlite3
import os
import json

SYSTEM_PROMPT = "You are a Text-to-SQL generator. Given a database schema and a natural language question, write a valid SQLite query that answers the question. Do not explain the query, just output the raw SQL."

def get_schema_ddl(db_path):
    con = sqlite3.connect(db_path)
    res = con.execute("SELECT sql FROM sqlite_master WHERE type='table' AND sql IS NOT NULL").fetchall()
    con.close()
    return "\n\n".join([row[0] for row in res])

def build_prompt(schema_ddl, question):
    return f"{SYSTEM_PROMPT}\n\nDatabase Schema:\n{schema_ddl}\n\nQuestion: {question}"

def format_chat(schema_ddl, question, query):
    prompt = build_prompt(schema_ddl, question)
    return {
        "messages": [
            {"role": "user", "content": prompt},
            {"role": "assistant", "content": query}
        ]
    }

def db_path(spider_root, db_id):
    return os.path.join(spider_root, "database", db_id, f"{db_id}.sqlite")

def generate_split(examples, spider_root, out_path):
    ddl_cache, n, skipped = {}, 0, 0
    with open(out_path, "w", encoding="utf-8") as f:
        for ex in examples:
            db_id = ex["db_id"]
            if db_id not in ddl_cache:
                try:
                    ddl_cache[db_id] = get_schema_ddl(db_path(spider_root, db_id))
                except Exception as e:
                    ddl_cache[db_id] = None
                    print(f"  WARN skip db {db_id}: {e}")
            if not ddl_cache[db_id]:
                skipped += 1
                continue
            rec = format_chat(ddl_cache[db_id], ex["question"], ex["query"])
            rec["db_id"] = db_id          # kept so the eval harness knows which DB to run against
            f.write(json.dumps(rec) + "\n")
            n += 1
    print(f"  wrote {n} to {out_path} (skipped {skipped})")
    return n

if __name__ == "__main__":
    from huggingface_hub import snapshot_download
    spider_root = os.path.join(
        snapshot_download("minktn/spider-data", repo_type="dataset"), "spider_data")
    out_dir = "data/spider"; os.makedirs(out_dir, exist_ok=True)
    train = (json.load(open(os.path.join(spider_root, "train_spider.json"), encoding="utf-8"))
             + json.load(open(os.path.join(spider_root, "train_others.json"), encoding="utf-8")))
    dev = json.load(open(os.path.join(spider_root, "dev.json"), encoding="utf-8"))
    nt = generate_split(train, spider_root, os.path.join(out_dir, "train.jsonl"))
    nd = generate_split(dev, spider_root, os.path.join(out_dir, "dev.jsonl"))
    assert nt >= 8000, f"train too small ({nt}); expected ~8600"
    assert nd == 1034, f"dev should be exactly 1034, got {nd}"
    print(f"OK train={nt} dev={nd}")