# Historical source snapshot. Comments about accuracy, data origin, or compliance
# are author claims, not verification; see training/README.md and REPRODUCIBILITY.md.
"""
Submission Packager and Validator for Sinhala MMLU Shared Task
Ensures predictions match Codabench format requirements:
- Valid JSON on each line
- Required keys: 'q_no', 'subject', 'category', 'level', 'answer'
- Answer must be non-empty string index ('1', '2', '3', '4', '5')
- Creates predictions.zip containing predictions.jsonl
"""

import json
import zipfile
import sys
from pathlib import Path


def validate_and_package(predictions_file: str, output_zip: str = None) -> bool:
    pred_path = Path(predictions_file)
    if not pred_path.exists():
        print(f"Error: {predictions_file} does not exist!")
        return False

    if output_zip is None:
        output_zip = pred_path.parent / "predictions.zip"
    else:
        output_zip = Path(output_zip)

    required_keys = {"q_no", "subject", "category", "level", "answer"}
    total_lines = 0
    valid_lines = 0
    errors = []

    with open(pred_path, "r", encoding="utf-8") as f:
        for idx, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue
            total_lines += 1
            try:
                record = json.loads(line)
            except json.JSONDecodeError as e:
                errors.append(f"Line {idx}: Malformed JSON ({e})")
                continue

            missing = required_keys - set(record.keys())
            if missing:
                errors.append(f"Line {idx}: Missing keys {missing}")
                continue

            ans = str(record["answer"]).strip()
            if not ans:
                errors.append(f"Line {idx}: Empty answer")
                continue

            valid_lines += 1

    print(f"Validated {valid_lines}/{total_lines} prediction records.")
    if errors:
        print(f"Found {len(errors)} validation errors:")
        for err in errors[:10]:
            print("  -", err)
        if len(errors) > 10:
            print(f"  ... and {len(errors) - 10} more errors.")
        return False

    print(f"Packing {pred_path} into {output_zip}...")
    with zipfile.ZipFile(output_zip, "w", zipfile.ZIP_DEFLATED) as z:
        # Note: Codabench requires the file inside the zip to be named predictions.jsonl
        z.write(pred_path, arcname="predictions.jsonl")

    print(f"SUCCESS: Submission ready at {output_zip}")
    print("Ready to upload to Codabench (up to 3 submissions/day).")
    return True


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python package_submission.py <path/to/predictions.jsonl> [path/to/output.zip]")
        sys.exit(1)

    input_file = sys.argv[1]
    out_file = sys.argv[2] if len(sys.argv) > 2 else None
    success = validate_and_package(input_file, out_file)
    sys.exit(0 if success else 1)