"""
Validates a trained model bundle directory before it's deployed to the API.

This deliberately mirrors the validation logic in
app/models_loader/registry.py so you catch a broken model bundle here,
during training/CI, rather than discovering it when the production API
refuses to start (or worse, if that fail-fast check ever has a gap).

Usage:
    python -m ml_pipeline.evaluation.validate_model_bundle --model-dir data/models/v1
"""

import argparse
import json
import sys
from pathlib import Path

import joblib

REQUIRED_FILES = {
    "sarimax_climate.pkl": "dict",
    "xgb_feasibility.pkl": "bundle",
    "xgb_yield.pkl": "bundle",
    "dt_trend.pkl": "bundle",
    "rf_recommendation.pkl": "bundle",
    "sarimax_market.pkl": "dict",
}


def validate(model_dir: Path) -> list[str]:
    issues: list[str] = []

    if not model_dir.exists():
        return [f"Model directory does not exist: {model_dir}"]

    for filename, expected_kind in REQUIRED_FILES.items():
        path = model_dir / filename
        if not path.exists():
            issues.append(f"MISSING: {filename}")
            continue

        try:
            obj = joblib.load(path)
        except Exception as exc:
            issues.append(f"CORRUPT: {filename} failed to load — {exc}")
            continue

        if expected_kind == "bundle":
            if not isinstance(obj, dict) or "model" not in obj:
                issues.append(
                    f"INVALID FORMAT: {filename} expected dict with 'model' key, "
                    f"got {type(obj)}"
                )
            elif "feature_cols" not in obj:
                issues.append(f"WARNING: {filename} has no 'feature_cols' key recorded")
        elif expected_kind == "dict":
            if not isinstance(obj, dict):
                issues.append(f"INVALID FORMAT: {filename} expected dict of fitted models, got {type(obj)}")
            elif len(obj) == 0:
                ack_path = model_dir / f"{filename}.empty_acknowledged"
                if ack_path.exists():
                    issues.append(
                        f"WARNING: {filename} is intentionally empty (acknowledged via "
                        f"{ack_path.name}) — predictions using this model will be degraded "
                        f"for every query until it is retrained with real data."
                    )
                else:
                    issues.append(
                        f"EMPTY: {filename} contains zero trained sub-models. If this is "
                        f"a deliberate placeholder for a known data gap (see "
                        f"docs/data_sourcing.md), create an empty marker file at "
                        f"{ack_path} to acknowledge this explicitly and downgrade this "
                        f"to a warning. Do not deploy an unacknowledged empty model."
                    )

        metadata_path = model_dir / f"{filename.replace('.pkl', '')}_metadata.json"
        if not metadata_path.exists():
            issues.append(f"WARNING: no metadata file for {filename} — cannot audit training provenance")
        else:
            with open(metadata_path, encoding="utf-8") as f:
                meta = json.load(f)
            if "trained_at_utc" not in meta:
                issues.append(f"WARNING: {metadata_path.name} missing 'trained_at_utc'")

    return issues


def main() -> None:
    parser = argparse.ArgumentParser(description="Validate a trained model bundle directory")
    parser.add_argument("--model-dir", required=True)
    args = parser.parse_args()

    issues = validate(Path(args.model_dir))

    errors = [i for i in issues if not i.startswith("WARNING")]
    warnings = [i for i in issues if i.startswith("WARNING")]

    if warnings:
        print("WARNINGS:")
        for w in warnings:
            print(f"  - {w}")

    if errors:
        print("\nVALIDATION FAILED — model bundle is NOT safe to deploy:")
        for e in errors:
            print(f"  - {e}")
        sys.exit(1)

    print(f"\nVALIDATION PASSED — {args.model_dir} is ready to deploy")


if __name__ == "__main__":
    main()