"""Audit master image role coverage before reimport / republish.

Examples:
  python -m app.jobs.audit_master_image_roles
  python -m app.jobs.audit_master_image_roles --limit 200
"""

from __future__ import annotations

import argparse
import json
from typing import Any, Dict, List, Optional, Sequence


def audit_master_image_roles(*, limit: Optional[int] = None, sample_limit: int = 25, session: Any = None) -> Dict[str, Any]:
    from sqlalchemy import select

    from db.master_product_images import COLLECTION_IMAGE_ROLES, PRODUCT_GALLERY_IMAGE_ROLES, SAMPLE_IMAGE_ROLES, normalize_image_role
    from db.models import CollectionCommerceProfile, MasterProduct, MasterProductImage
    from db.tribeca_sku_parse import canonical_collection_key, collection_asset_key, collection_name_from_sku

    def _run(active_session: Any) -> Dict[str, Any]:
        products = list(active_session.scalars(select(MasterProduct).where(MasterProduct.is_active.is_(True)).order_by(MasterProduct.sku)).all())
        if limit:
            products[:] = products[: int(limit)]
        product_by_sku = {str(product.sku): product for product in products}
        wanted_skus = set(product_by_sku)
        images = list(
            active_session.scalars(
                select(MasterProductImage).where(MasterProductImage.sku.in_(sorted(wanted_skus)))
            ).all()
        )
        profiles = list(active_session.scalars(select(CollectionCommerceProfile).where(CollectionCommerceProfile.is_active.is_(True))).all())

        images_by_sku: Dict[str, List[Any]] = {}
        role_counts: Dict[str, int] = {}
        rta_collection_asset_count = 0
        collection_mismatch_count = 0
        for image in images:
            images_by_sku.setdefault(str(image.sku), []).append(image)
            role = normalize_image_role(getattr(image, "image_role", None)) or "missing"
            role_counts[role] = role_counts.get(role, 0) + 1

        missing_images: List[str] = []
        missing_product_shot: List[str] = []
        rta_collection_assets: List[Dict[str, Any]] = []
        collection_mismatches: List[Dict[str, Any]] = []

        for sku, product in product_by_sku.items():
            rows = images_by_sku.get(sku) or []
            if not rows:
                missing_images.append(sku)
                continue
            normalized_roles = {
                normalize_image_role(getattr(row, "image_role", None)) or "missing"
                for row in rows
            }
            is_sample = sku.upper().endswith("-SD") or "-SDP" in sku.upper()
            if not is_sample and not (normalized_roles & PRODUCT_GALLERY_IMAGE_ROLES):
                if not normalized_roles & SAMPLE_IMAGE_ROLES:
                    missing_product_shot.append(sku)
            for row in rows:
                role = normalize_image_role(getattr(row, "image_role", None)) or "missing"
                if sku.upper().startswith("RTA-") and role in COLLECTION_IMAGE_ROLES:
                    rta_collection_asset_count += 1
                    if len(rta_collection_assets) < sample_limit:
                        rta_collection_assets.append(
                            {
                                "sku": sku,
                                "file_name": row.file_name,
                                "image_role": role,
                            }
                        )
                expected_collection = collection_asset_key(str(row.file_name or ""))
                canonical_collection = collection_name_from_sku(sku) or getattr(product, "collection", "")
                actual_collection = canonical_collection_key(str(canonical_collection or ""))
                if expected_collection and role in COLLECTION_IMAGE_ROLES and expected_collection != actual_collection:
                    collection_mismatch_count += 1
                    if len(collection_mismatches) < sample_limit:
                        collection_mismatches.append(
                            {
                                "sku": sku,
                                "collection": getattr(product, "collection", None),
                                "file_name": row.file_name,
                                "expected_collection_key": expected_collection,
                                "image_role": role,
                            }
                        )

        sample_skus_missing_images: List[Dict[str, Any]] = []
        sample_skus_missing_sample_roles: List[Dict[str, Any]] = []
        for profile in profiles:
            sample_sku = str(getattr(profile, "sample_master_sku", "") or "").strip()
            if not sample_sku:
                continue
            rows = images_by_sku.get(sample_sku) or []
            if not rows:
                if len(sample_skus_missing_images) < sample_limit:
                    sample_skus_missing_images.append(
                        {
                            "collection_name": profile.collection_name,
                            "sample_master_sku": sample_sku,
                        }
                    )
                continue
            normalized_roles = {
                normalize_image_role(getattr(row, "image_role", None)) or "missing"
                for row in rows
            }
            if not normalized_roles & SAMPLE_IMAGE_ROLES:
                if len(sample_skus_missing_sample_roles) < sample_limit:
                    sample_skus_missing_sample_roles.append(
                        {
                            "collection_name": profile.collection_name,
                            "sample_master_sku": sample_sku,
                            "roles": sorted(normalized_roles),
                        }
                    )

        return {
            "status": "ok",
            "product_count": len(product_by_sku),
            "image_count": len(images),
            "role_counts": dict(sorted(role_counts.items())),
            "missing_images_count": len(missing_images),
            "sample_missing_images_count": len(sample_skus_missing_images),
            "sample_missing_sample_roles_count": len(sample_skus_missing_sample_roles),
            "missing_product_shot_count": len(missing_product_shot),
            "rta_collection_asset_count": rta_collection_asset_count,
            "collection_mismatch_count": collection_mismatch_count,
            "sample_missing_images": sample_skus_missing_images,
            "sample_missing_sample_roles": sample_skus_missing_sample_roles,
            "missing_images_sample": missing_images[:sample_limit],
            "missing_product_shot_sample": missing_product_shot[:sample_limit],
            "rta_collection_assets_sample": rta_collection_assets,
            "collection_mismatches_sample": collection_mismatches,
        }

    if session is not None:
        return _run(session)

    from db.session import get_session

    with get_session() as managed_session:
        return _run(managed_session)


def main(argv: Optional[Sequence[str]] = None) -> int:
    parser = argparse.ArgumentParser(description="Audit master image role coverage")
    parser.add_argument("--limit", type=int, default=None, help="Limit active SKUs for a quick sample")
    parser.add_argument("--sample-limit", type=int, default=25, help="Sample size per issue bucket")
    args = parser.parse_args(list(argv) if argv is not None else None)
    print(json.dumps(audit_master_image_roles(limit=args.limit, sample_limit=args.sample_limit), indent=2))
    return 0


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