# convert.py
import json
import os
from datasets import load_dataset

os.makedirs("datasets/msmarco/qrels", exist_ok=True)

N_DOCS = 100000

# Convert only 100K corpus docs
print("Converting corpus (100K only)...")
corpus = load_dataset("parquet", data_files="datasets/msmarco/corpus/*.parquet")["train"]
with open("datasets/msmarco/corpus.jsonl", "w", encoding="utf-8") as f:
    for i, row in enumerate(corpus):
        if i >= N_DOCS:
            break
        f.write(json.dumps({
            "_id": str(row["_id"]),
            "title": row.get("title", ""),
            "text": row["text"]
        }) + "\n")
print("✅ corpus.jsonl done (100K)")

# Convert queries
print("Converting queries...")
queries = load_dataset("parquet", data_files="datasets/msmarco/queries/*.parquet")["train"]
with open("datasets/msmarco/queries.jsonl", "w", encoding="utf-8") as f:
    for row in queries:
        f.write(json.dumps({
            "_id": str(row["_id"]),
            "text": row["text"]
        }) + "\n")
print("✅ queries.jsonl done")

# Download qrels
print("Downloading qrels...")
qrels = load_dataset("BeIR/msmarco-qrels", split="validation")
with open("datasets/msmarco/qrels/dev.tsv", "w", encoding="utf-8") as f:
    f.write("query-id\tcorpus-id\tscore\n")
    for row in qrels:
        f.write(f"{row['query-id']}\t{row['corpus-id']}\t{row['score']}\n")
print("✅ qrels/dev.tsv done")

print("\n✅ All done! Now run main.py")