from __future__ import annotations

import argparse
import sys
from pathlib import Path

from common import iter_model_card_paths, read_yaml_mapping

REQUIRED_FIELDS = {
    "id",
    "name",
    "type",
    "task",
    "language",
    "source",
    "license",
    "status",
    "storage_path",
    "deployment",
}

VALID_STATUSES = {"candidate", "draft", "staging", "production", "archived"}


def load_yaml(path: Path) -> dict:
    return read_yaml_mapping(path)


def validate_card(path: Path) -> list[str]:
    errors: list[str] = []

    try:
        data = load_yaml(path)
    except Exception as exc:
        return [f"{path}: cannot parse YAML: {exc}"]

    missing = sorted(REQUIRED_FIELDS - set(data))
    for field in missing:
        errors.append(f"{path}: missing required field `{field}`")

    status = data.get("status")
    if status is not None and status not in VALID_STATUSES:
        errors.append(
            f"{path}: invalid status `{status}`, expected one of {sorted(VALID_STATUSES)}"
        )

    for field in ("task", "language"):
        value = data.get(field)
        if value is not None and not isinstance(value, list):
            errors.append(f"{path}: `{field}` must be a list")

    deployment = data.get("deployment")
    if deployment is not None and not isinstance(deployment, dict):
        errors.append(f"{path}: `deployment` must be a mapping")

    return errors


def main() -> int:
    parser = argparse.ArgumentParser(description="Validate model cards.")
    parser.add_argument("--include-examples", action="store_true")
    args = parser.parse_args()

    cards = iter_model_card_paths(include_examples=args.include_examples)
    if not cards:
        print("No model cards found.")
        return 0

    errors: list[str] = []
    for card in cards:
        errors.extend(validate_card(card))

    if errors:
        print("Model card validation failed:")
        for error in errors:
            print(f"- {error}")
        return 1

    print(f"Validated {len(cards)} model card(s).")
    return 0


if __name__ == "__main__":
    sys.exit(main())