"""Sinh toan bo file du lieu cho Space tu du an goc.

    python build_data.py "/duong/dan/Thom Paper"

Ghi ra space/data/:
    dataset.json   10.961 cap hoi dap (train + dev + test)
    facets.json    gia tri cho cac bo loc
    lessons.json   cay lop > chu de > bai hoc
    metrics.json   chi so thuc nghiem + bao cao ro ri du lieu
    testset.json   1.073 cau test kem du doan cua 6 mo hinh
    errors.json    vi du phan tich loi, nhom theo kieu
    recorded.jsonl du doan da ghi — dung khi chua co checkpoint fine-tune
    built.json     thoi diem build
"""

import csv
import json
import re
import sys
from datetime import datetime, timezone
from pathlib import Path

HERE = Path(__file__).parent
OUT = HERE / "data"

MODELS = [
    ("phobert", "PhoBERT", "extractive"),
    ("mbert", "mBERT", "extractive"),
    ("xlmr", "XLM-RoBERTa", "extractive"),
    ("vit5", "ViT5", "generative"),
    ("bartpho", "BARTpho", "generative"),
    ("gpt", "GPT", "generative"),
]
LABELS = [label for _, label, _ in MODELS]

NEEDED = [
    "data/processed/train.csv",
    "data/processed/dev.csv",
    "data/processed/test.csv",
    "results/final_results.json",
    "results/results_by_grade.csv",
    "results/du_doan_6_mo_hinh.csv",
    "results/vi_du_phan_tich_loi.csv",
    "results/data_split_report.json",
]


def desegment(value):
    """Du lieu PhoBERT duoc tach tu nen co dau gach duoi — bo di khi hien thi."""
    return re.sub(r"\s+", " ", str(value or "").replace("_", " ")).strip()


def num(value):
    if value in ("", None):
        return None
    try:
        return float(value)
    except ValueError:
        return None


def read_csv(root: Path, rel: str):
    with (root / rel).open(encoding="utf-8-sig", newline="") as fh:
        return list(csv.DictReader(fh))


def read_json(root: Path, rel: str):
    return json.loads((root / rel).read_text(encoding="utf-8"))


def write(name: str, value) -> None:
    path = OUT / name
    path.write_text(json.dumps(value, ensure_ascii=False), encoding="utf-8")
    print(f"  {name:<16} {path.stat().st_size / 1024:,.0f} KB")


# ----------------------------------------------------------------- dataset
def build_dataset(root: Path):
    items = []
    for split in ("train", "dev", "test"):
        for r in read_csv(root, f"data/processed/{split}.csv"):
            items.append(
                {
                    "id": r["id"],
                    "grade": int(r["grade"]),
                    "topic": r["topic"],
                    "lesson": r["lesson"],
                    "page": r["page"],
                    "section": r["section"],
                    "contextId": r["context_id"],
                    "context": r["context"],
                    "question": r["question"],
                    "answer": r["answer"],
                    "answerStart": int(r["answer_start"]) if r["answer_start"] else None,
                    "answerEnd": int(r["answer_end"]) if r["answer_end"] else None,
                    "intent": r["intent"],
                    "difficulty": r["difficulty"],
                    "questionType": r["question_type"],
                    "source": r["source"],
                    "split": split,
                }
            )
    return items


def build_lessons(items):
    grades = {}
    for it in items:
        topics = grades.setdefault(it["grade"], {})
        lessons = topics.setdefault(it["topic"], {})
        entry = lessons.setdefault(
            it["lesson"], {"lesson": it["lesson"], "page": it["page"], "questions": 0, "contexts": set()}
        )
        entry["questions"] += 1
        entry["contexts"].add(it["contextId"])

    out = []
    for grade in sorted(grades):
        topics = []
        for topic in sorted(grades[grade]):
            lessons = sorted(grades[grade][topic].values(), key=lambda l: l["lesson"])
            topics.append(
                {
                    "topic": topic,
                    "lessons": [
                        {"lesson": l["lesson"], "page": l["page"], "questions": l["questions"], "contexts": len(l["contexts"])}
                        for l in lessons
                    ],
                }
            )
        out.append({"grade": grade, "topics": topics})
    return out


def build_facets(items):
    uniq = lambda key: sorted({it[key] for it in items if it[key]})
    return {
        "grades": sorted({it["grade"] for it in items}),
        "topics": uniq("topic"),
        "difficulties": uniq("difficulty"),
        "questionTypes": uniq("questionType"),
        "splits": ["train", "dev", "test"],
        "total": len(items),
    }


def build_stats(items):
    def count(key):
        out = {}
        for it in items:
            out[str(it[key])] = out.get(str(it[key]), 0) + 1
        return out

    contexts = {it["contextId"]: it["context"] for it in items}
    return {
        "total": len(items),
        "contexts": len(contexts),
        "byGrade": count("grade"),
        "bySplit": count("split"),
        "byDifficulty": count("difficulty"),
        "byQuestionType": count("questionType"),
        "answerLengthAvg": round(sum(len(it["answer"] or "") for it in items) / len(items), 1),
        "contextLengthAvg": round(sum(len(c) for c in contexts.values()) / len(contexts)),
    }


# ------------------------------------------------------------------ metrics
def build_metrics(root: Path):
    final = read_json(root, "results/final_results.json")
    split = read_json(root, "results/data_split_report.json")

    models = []
    for key, label, kind in MODELS:
        block = final.get("extractive", {}).get(label) or final.get("generative", {}).get(label) or {}
        models.append(
            {
                "key": key,
                "label": label,
                "kind": kind,
                "em": round(block["EM"], 2) if "EM" in block else None,
                "f1": round(block["F1"], 2) if "F1" in block else None,
                "bleu": round(block["BLEU"], 2) if "BLEU" in block else None,
                "rougeL": round(block["ROUGE-L"], 2) if "ROUGE-L" in block else None,
            }
        )

    per_grade = [
        {
            "label": r["Mô hình"],
            "grade": int(r["Lớp"]),
            "em": num(r["EM"]),
            "f1": num(r["F1"]),
            "bleu": num(r["BLEU"]),
            "rougeL": num(r["ROUGE-L"]),
        }
        for r in read_csv(root, "results/results_by_grade.csv")
    ]

    return {
        "models": models,
        "perGrade": per_grade,
        "counts": {"train": final["n_train"], "dev": final["n_dev"], "test": final["n_test"]},
        "split": {
            "rowsRaw": split["rows_raw"],
            "rowsClean": split["rows_after_cleaning"],
            "contextClusters": split["context_clusters"],
            "uniqueContexts": split["unique_context_texts"],
            "seed": split["seed"],
            "jaccard": split["jaccard_threshold"],
            "leakage": split["leakage_checks"],
        },
    }


# ---------------------------------------------- testset + du doan da ghi
def build_testset_and_recorded(root: Path):
    rows = read_csv(root, "results/du_doan_6_mo_hinh.csv")
    testset = []
    recorded = []
    for r in rows:
        preds = {}
        for key, label, _ in MODELS:
            preds[key] = {
                "answer": desegment(r.get(label)),
                "em": num(r.get(f"{label}_em")),
                "f1": num(r.get(f"{label}_f1")),
            }
        testset.append(
            {
                "id": r["id"],
                "grade": int(r["grade"]),
                "topic": r.get("topic", ""),
                "lesson": r.get("lesson", ""),
                "question": r["question"],
                "context": r["context"],
                "gold": desegment(r["gold"]),
                "preds": preds,
                "extCorrect": num(r.get("n_ext_dung")),
                "genF1Avg": num(r.get("gen_f1_tb")),
            }
        )
        recorded.append(
            {
                "id": r["id"],
                "grade": r["grade"],
                "question": r["question"],
                "context": r["context"],
                "gold": r["gold"],
                "lesson": r.get("lesson", ""),
                "topic": r.get("topic", ""),
                "predictions": {label: (r.get(label) or "").strip() for label in LABELS},
                "f1": {label: round(float(r[f"{label}_f1"]), 4) for label in LABELS if num(r.get(f"{label}_f1")) is not None},
            }
        )
    return testset, recorded


def build_errors(root: Path):
    groups = {}
    for r in read_csv(root, "results/vi_du_phan_tich_loi.csv"):
        name = r["nhóm"]
        group = groups.setdefault(name, {"group": name, "items": []})
        group["items"].append(
            {
                "id": r["id"],
                "grade": int(r["lớp"]),
                "lesson": r["bài"],
                "question": r["câu hỏi"],
                "gold": desegment(r["đáp án"]),
                "context": r["ngữ cảnh"],
                "preds": {
                    key: {"answer": desegment(r.get(label)), "f1": num(r.get(f"{label}_F1"))}
                    for key, label, _ in MODELS
                },
            }
        )
    return list(groups.values())


# -------------------------------------------------------------------- main
def main() -> int:
    root = Path(sys.argv[1] if len(sys.argv) > 1 else HERE.parent.parent).expanduser().resolve()
    missing = [f for f in NEEDED if not (root / f).exists()]
    if missing:
        print(f"Khong tim thay trong {root}:")
        for f in missing:
            print(f"  - {f}")
        print('\nChay lai voi duong dan goc, vi du:\n  python build_data.py "/Volumes/MyData/paper/Thom Paper"')
        return 1

    OUT.mkdir(exist_ok=True)
    print(f"Doc tu: {root}")

    items = build_dataset(root)
    testset, recorded = build_testset_and_recorded(root)

    write("dataset.json", items)
    write("facets.json", build_facets(items))
    write("stats.json", build_stats(items))
    write("lessons.json", build_lessons(items))
    write("metrics.json", build_metrics(root))
    write("testset.json", testset)
    write("errors.json", build_errors(root))

    path = OUT / "recorded.jsonl"
    with path.open("w", encoding="utf-8") as fh:
        for row in recorded:
            fh.write(json.dumps(row, ensure_ascii=False) + "\n")
    print(f"  {'recorded.jsonl':<16} {path.stat().st_size / 1024:,.0f} KB")

    write(
        "built.json",
        {"builtAt": datetime.now(timezone.utc).isoformat(), "source": str(root), "items": len(items)},
    )
    print(f"\nXong — {len(items):,} cap hoi dap, {len(testset):,} cau test.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())