from __future__ import annotations

import argparse
import datetime as dt
import json
from pathlib import Path
from typing import Any

from common import iter_model_card_paths, localize_workspace_path, read_yaml_mapping


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_REPORT = ROOT / "reports" / "model-storage.json"
REQUIRED_MODEL_STATUSES = {"staging", "production"}


def has_any(path: Path, patterns: list[str]) -> bool:
    return any(path.glob(pattern) for pattern in patterns)


def has_weight_file(path: Path) -> bool:
    for file_path in path.iterdir() if path.exists() else []:
        if not file_path.is_file():
            continue
        name = file_path.name.lower()
        if name.endswith((".safetensors", ".bin", ".gguf")):
            return True
    return False


def check_gguf(path: Path, card: dict[str, Any]) -> dict[str, Any]:
    filename = card.get("filename")
    if not filename:
        return {"status": "failed", "reason": "filename is missing in model card"}
    file_path = path / str(filename)
    if not file_path.exists():
        return {"status": "missing", "reason": f"file is missing: {file_path.relative_to(ROOT)}"}
    size = file_path.stat().st_size
    expected_size = card.get("file_size_bytes")
    if expected_size and size != int(expected_size):
        return {
            "status": "partial",
            "reason": f"size mismatch: {size} != {expected_size}",
            "size_bytes": size,
            "expected_size_bytes": expected_size,
        }
    return {"status": "ok", "size_bytes": size, "expected_size_bytes": expected_size}


def check_hf_model(path: Path) -> dict[str, Any]:
    required = ["config.json"]
    missing = [name for name in required if not (path / name).exists()]
    index_path = path / "model.safetensors.index.json"
    missing_shards: list[str] = []
    if index_path.exists():
        try:
            index = json.loads(index_path.read_text(encoding="utf-8"))
            shard_names = sorted(set((index.get("weight_map") or {}).values()))
            missing_shards = [name for name in shard_names if not (path / name).exists()]
        except json.JSONDecodeError:
            return {"status": "failed", "reason": "model.safetensors.index.json is invalid"}
    has_weights = has_weight_file(path)
    has_tokenizer = has_any(path, ["tokenizer.json", "tokenizer.model", "vocab.json"])
    if missing:
        return {"status": "missing", "reason": f"missing required file(s): {', '.join(missing)}"}
    if missing_shards:
        preview = ", ".join(missing_shards[:4])
        suffix = f" and {len(missing_shards) - 4} more" if len(missing_shards) > 4 else ""
        return {"status": "partial", "reason": f"missing shard file(s): {preview}{suffix}"}
    if not has_weights:
        return {"status": "metadata-only", "reason": "model weights are missing"}
    if not has_tokenizer:
        return {"status": "partial", "reason": "tokenizer files are missing"}
    return {"status": "ok"}


def check_diffusers_model(path: Path) -> dict[str, Any]:
    missing = [name for name in ["model_index.json"] if not (path / name).exists()]
    if missing:
        return {"status": "missing", "reason": f"missing required file(s): {', '.join(missing)}"}
    if not any(path.rglob("*.safetensors")) and not any(path.rglob("*.bin")):
        return {"status": "metadata-only", "reason": "diffusers weights are missing"}
    return {"status": "ok"}


def check_adapter(path: Path) -> dict[str, Any]:
    if not path.exists():
        return {"status": "missing", "reason": f"adapter path is missing: {path}"}
    if has_any(path, ["adapter_config.json", "*.safetensors", "*.bin"]):
        return {"status": "ok"}
    return {"status": "metadata-only", "reason": "adapter artifact files are missing"}


def resolve_storage_path(raw_path: str, models_root: Path | None = None) -> Path:
    if models_root and raw_path.startswith("/models/"):
        return models_root / raw_path.removeprefix("/models/")
    return localize_workspace_path(raw_path)


def check_card(path: Path, *, models_root: Path | None = None) -> dict[str, Any]:
    card = read_yaml_mapping(path)
    storage_path = resolve_storage_path(str(card.get("storage_path") or ""), models_root=models_root)
    model_status = str(card.get("status") or "draft")
    item = {
        "id": card.get("id"),
        "name": card.get("name"),
        "type": card.get("type"),
        "model_status": model_status,
        "format": card.get("format"),
        "quantization": card.get("quantization"),
        "runtime": (card.get("deployment") or {}).get("runtime"),
        "served_model_name": (card.get("deployment") or {}).get("served_model_name"),
        "filename": card.get("filename"),
        "storage_path": str(storage_path),
        "card_path": str(path.relative_to(ROOT)),
        "status": "missing",
        "reason": None,
        "required": model_status in REQUIRED_MODEL_STATUSES,
    }
    if not storage_path.exists():
        item["reason"] = "storage path is missing"
        return item

    model_format = str(card.get("format") or "").lower()
    model_type = str(card.get("type") or "").lower()
    if model_format == "gguf":
        result = check_gguf(storage_path, card)
    elif model_format == "diffusers" or model_type == "image-diffusion-model":
        result = check_diffusers_model(storage_path)
    elif model_type == "lora-adapter":
        result = check_adapter(storage_path)
    else:
        result = check_hf_model(storage_path)
    item.update(result)
    return item


def main() -> int:
    parser = argparse.ArgumentParser(description="Check local model storage against model cards.")
    parser.add_argument("--report", type=Path, default=DEFAULT_REPORT)
    parser.add_argument("--print", action="store_true")
    parser.add_argument("--no-report", action="store_true")
    parser.add_argument("--warn-only", action="store_true", help="Return success even when models are partial or missing.")
    parser.add_argument("--strict", action="store_true", help="Fail on incomplete draft/candidate models too.")
    parser.add_argument(
        "--models-root",
        type=Path,
        default=None,
        help="Override /models paths, for example Z:/LLM/models or /models inside the GPU host container.",
    )
    args = parser.parse_args()

    models_root = args.models_root.resolve() if args.models_root else None
    models = [check_card(path, models_root=models_root) for path in iter_model_card_paths()]
    incomplete_statuses = {"failed", "missing", "partial"}
    incomplete = [model["id"] for model in models if model["status"] in incomplete_statuses]
    required_failed = [
        model["id"]
        for model in models
        if model["required"] and model["status"] in incomplete_statuses
    ]
    failed = incomplete if args.strict else required_failed
    report = {
        "created_at": dt.datetime.now(dt.UTC).isoformat(),
        "models_root": str(models_root) if models_root else None,
        "status": "failed" if failed else "ok",
        "failed": failed,
        "required_failed": required_failed,
        "planned_incomplete": [
            model["id"]
            for model in models
            if not model["required"] and model["status"] in incomplete_statuses
        ],
        "strict": args.strict,
        "counts": {
            "ok": sum(1 for model in models if model["status"] == "ok"),
            "missing": sum(1 for model in models if model["status"] == "missing"),
            "partial": sum(1 for model in models if model["status"] == "partial"),
            "metadata_only": sum(1 for model in models if model["status"] == "metadata-only"),
            "failed": sum(1 for model in models if model["status"] == "failed"),
        },
        "models": models,
    }
    if not args.no_report:
        args.report.parent.mkdir(parents=True, exist_ok=True)
        args.report.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    if args.print:
        print(json.dumps(report, ensure_ascii=False, indent=2))
    else:
        print(f"Model storage status: {report['status']}")
    return 0 if args.warn_only else 1 if failed else 0


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