#!/usr/bin/env python3
"""Validate generated Potato data files."""

from __future__ import annotations

import argparse
import csv
import json
import sys
from pathlib import Path


PROJECT_DIR = Path(__file__).resolve().parents[1]
DEFAULT_DATA_DIR = PROJECT_DIR / "my-annotation-task" / "data"

DIMENSIONS = [
    "Content Correctness",
    "Learner-State Assessment",
    "Issue Localization",
    "Disclosure Appropriateness",
    "Providing Guidance",
    "Coherence",
    "Actionability",
    "Clarity",
    "Conciseness",
    "Humanness",
]

VALID_LABELS = {"Yes", "To some extent", "No"}
REQUIRED_COLUMNS = [
    "id",
    "split",
    "dialog_context",
    "correct_solution",
    "tutor_response",
    "text2show_html",
    *DIMENSIONS,
]


def read_csv(path: Path) -> list[dict[str, str]]:
    with path.open(newline="", encoding="utf-8") as handle:
        return list(csv.DictReader(handle))


def validate_csv(path: Path, expected_split: str | None = None) -> int:
    rows = read_csv(path)
    if not rows:
        raise ValueError(f"{path.name}: no rows")
    missing = [column for column in REQUIRED_COLUMNS if column not in rows[0]]
    if missing:
        raise ValueError(f"{path.name}: missing columns {missing}")

    ids = set()
    for row_number, row in enumerate(rows, start=2):
        item_id = row["id"]
        if not item_id:
            raise ValueError(f"{path.name}:{row_number}: empty id")
        if item_id in ids:
            raise ValueError(f"{path.name}:{row_number}: duplicate id {item_id}")
        ids.add(item_id)
        if expected_split and row["split"] != expected_split:
            raise ValueError(f"{path.name}:{row_number}: expected split {expected_split}, got {row['split']}")
        for column in REQUIRED_COLUMNS:
            if not row[column].strip():
                raise ValueError(f"{path.name}:{row_number}: empty {column}")
        for dimension in DIMENSIONS:
            if row[dimension] not in VALID_LABELS:
                raise ValueError(f"{path.name}:{row_number}: invalid {dimension}: {row[dimension]!r}")
    return len(rows)


def validate_json(path: Path, key_name: str) -> int:
    payload = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(payload, list) or not payload:
        raise ValueError(f"{path.name}: expected a non-empty list")
    ids = set()
    for index, item in enumerate(payload, start=1):
        item_id = item.get("id")
        if not item_id:
            raise ValueError(f"{path.name}:{index}: empty id")
        if item_id in ids:
            raise ValueError(f"{path.name}:{index}: duplicate id {item_id}")
        ids.add(item_id)
        labels = item.get(key_name)
        if not isinstance(labels, dict):
            raise ValueError(f"{path.name}:{index}: missing {key_name}")
        for dimension in DIMENSIONS:
            if labels.get(dimension) not in VALID_LABELS:
                raise ValueError(f"{path.name}:{index}: invalid {dimension}: {labels.get(dimension)!r}")
    return len(payload)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--data-dir", type=Path, default=DEFAULT_DATA_DIR)
    args = parser.parse_args()

    checks = {
        "train_csv": validate_csv(args.data_dir / "math_annotator_training_set_with_id_text2show.csv", "train"),
        "test_csv": validate_csv(args.data_dir / "math_annotator_testing_set_with_id_text2show.csv", "test"),
        "combined_csv": validate_csv(args.data_dir / "math_annotator_demo_all_with_id_text2show.csv"),
        "training_questions": validate_json(args.data_dir / "training_questions.json", "correct_answers"),
        "gold_standards": validate_json(args.data_dir / "gold_standards.json", "gold_label"),
    }
    print(json.dumps(checks, indent=2))
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except Exception as exc:
        print(f"Validation failed: {exc}", file=sys.stderr)
        raise SystemExit(1)