normalization/merger.py

"""
Report merger — turns a dict of {provider_name: ProviderResult} into a
UnifiedFaceReport.

The merger:
  1. Routes each ProviderResult to the correct collector based on capability.
  2. Preserves every ProviderResult as Evidence (raw + normalized + timestamp).
  3. Delegates confidence scoring to confidence.engine.
  4. Delegates conflict detection to confidence.conflicts.
  5. Attaches limitations collected from providers + system-wide.
"""

from __future__ import annotations

from typing import Dict, List

from confidence.engine import ConfidenceEngine
from confidence.conflicts import ConflictDetector
from models.reports import (
    UnifiedFaceReport,
    FaceDetection,
    FaceMatch,
    ImageAnalysisResult,
    MetadataResult,
    ForensicsResult,
    OCRResult,
    ObjectDetectionResult,
    SceneResult,
    NSFWResult,
    AIDetectionResult,
    EmbeddingResult,
    Evidence,
    ReportMetadata,
)
from normalization.schema import (
    NormalizedBox,
    NormalizedMatch,
    NormalizedScrapeImage,
    NormalizedReverseMatch,
    NormalizedImageAnalysis,
    NormalizedMetadata,
    NormalizedForensics,
)
from pipeline.postprocessing import ResultPostprocessor
from providers.base import ProviderResult


class ReportMerger:
    """Merges provider results into a UnifiedFaceReport."""

    def __init__(
        self,
        confidence_engine: ConfidenceEngine,
        conflict_detector: ConflictDetector,
    ) -> None:
        self._confidence = confidence_engine
        self._conflicts = conflict_detector

    def merge(
        self,
        results: Dict[str, ProviderResult],
        image_hash: str,
        job_id: str,
        total_elapsed_ms: float,
        kind: str = "detection",
    ) -> UnifiedFaceReport:
        boxes = self._collect_boxes(results)
        matches = self._collect_matches(results)
        scraped = self._collect_scraped(results)
        reverse_matches = self._collect_reverse(results)
        image_analyses = self._collect_image_analyses(results)
        metadata_extractions = self._collect_metadata(results)
        forensics = self._collect_forensics(results)
        evidence = self._build_evidence(results)

        # Build FaceDetection objects with confidence
        detections: List[FaceDetection] = []
        for i, nbox in enumerate(boxes):
            score = self._confidence.score_detection(nbox, results)
            detections.append(FaceDetection(
                box={"x": nbox.x, "y": nbox.y, "w": nbox.w, "h": nbox.h},
                confidence=score,
                landmarks=nbox.landmarks,
                detected_by=[nbox.detector],
            ))

        # Build FaceMatch objects
        face_matches: List[FaceMatch] = []
        for nm in matches:
            score = self._confidence.score_match(nm)
            face_matches.append(FaceMatch(
                query_face_index=nm.query_face_index,
                best_match=nm.best_match,
                confidence=score,
                distances=nm.distances,
            ))

        # Build ImageAnalysisResult objects
        image_analysis_models: List[ImageAnalysisResult] = []
        for nia in image_analyses:
            score = self._confidence.score_image_analysis(nia)
            image_analysis_models.append(ImageAnalysisResult(
                provider=nia.provider,
                quality_score=nia.quality_score,
                brightness=nia.brightness,
                contrast=nia.contrast,
                sharpness=nia.sharpness,
                noise_level=nia.noise_level,
                width=nia.width,
                height=nia.height,
                channels=nia.channels,
                color_profile=nia.color_profile,
                dominant_colors=nia.dominant_colors,
                aspects=nia.aspects,
                confidence=score,
            ))

        # Build MetadataResult objects
        metadata_models: List[MetadataResult] = []
        for nm in metadata_extractions:
            score = self._confidence.score_metadata(nm)
            metadata_models.append(MetadataResult(
                provider=nm.provider,
                format=nm.format,
                exif=nm.exif,
                xmp=nm.xmp,
                iptc=nm.iptc,
                gps=nm.gps,
                camera_make=nm.camera_make,
                camera_model=nm.camera_model,
                software=nm.software,
                capture_time=nm.capture_time,
                confidence=score,
            ))

        # Build ForensicsResult objects
        forensics_models: List[ForensicsResult] = []
        for nf in forensics:
            score = self._confidence.score_forensics(nf)
            forensics_models.append(ForensicsResult(
                provider=nf.provider,
                integrity_score=nf.integrity_score,
                is_duplicate=nf.is_duplicate,
                duplicate_of=nf.duplicate_of,
                similarity_score=nf.similarity_score,
                manipulation_indicators=nf.manipulation_indicators,
                ela_score=nf.ela_score,
                noise_inconsistency=nf.noise_inconsistency,
                details=nf.details,
                confidence=score,
            ))

        # Detect cross-provider conflicts
        conflicts = self._conflicts.detect(results, boxes, matches)

        # Aggregate limitations
        limitations: List[str] = []
        for r in results.values():
            if not r.success and r.error:
                limitations.append(f"{r.provider}: {r.error}")
        if not any(r.success for r in results.values()) and results:
            limitations.append("All providers failed for this job")

        # Collect new capability results
        ocr_results = self._collect_ocr(results)
        object_detections = self._collect_object_detections(results)
        scene_results = self._collect_scene(results)
        nsfw_results = self._collect_nsfw(results)
        ai_detection_results = self._collect_ai_detection(results)
        embedding_results = self._collect_embeddings(results)

        # Metadata
        succeeded = [r.provider for r in results.values() if r.success]
        failed = [r.provider for r in results.values() if not r.success]
        metadata = ReportMetadata(
            job_id=job_id,
            image_hash=image_hash,
            total_elapsed_ms=total_elapsed_ms,
            providers_invoked=list(results.keys()),
            providers_succeeded=succeeded,
            providers_failed=failed,
            limitations=limitations,
        )

        report = UnifiedFaceReport(
            metadata=metadata,
            detections=detections,
            matches=face_matches,
            scraped_images=[s.__dict__ for s in scraped],
            reverse_matches=[r.__dict__ for r in reverse_matches],
            image_analyses=image_analysis_models,
            metadata_extractions=metadata_models,
            forensics=forensics_models,
            ocr_results=ocr_results,
            object_detections=object_detections,
            scene_results=scene_results,
            nsfw_results=nsfw_results,
            ai_detection_results=ai_detection_results,
            embedding_results=embedding_results,
            evidence=evidence,
            conflicts=conflicts,
            overall_confidence=self._confidence.score_overall(
                report_detections=detections,
                matches=face_matches,
                conflicts=conflicts,
                image_analyses=image_analysis_models,
                forensics=forensics_models,
            ),
        )
        return report

    # ------------------------------------------------------------------ #
    # Collectors — one per capability
    # ------------------------------------------------------------------ #
    def _collect_boxes(self, results: Dict[str, ProviderResult]) -> List[NormalizedBox]:
        out: List[NormalizedBox] = []
        for r in results.values():
            if not r.success or r.capability.value != "detection":
                continue
            norm = r.normalized
            boxes = norm.get("boxes", [])
            confs = norm.get("confidences", [1.0] * len(boxes))
            lms = norm.get("landmarks")
            for i, b in enumerate(boxes):
                lm = lms[i] if isinstance(lms, list) and i < len(lms) else (lms if isinstance(lms, dict) else None)
                out.append(NormalizedBox(
                    x=b["x"], y=b["y"], w=b["w"], h=b["h"],
                    confidence=float(confs[i]) if i < len(confs) else 1.0,
                    detector=r.provider,
                    landmarks=lm,
                ))
        return out

    def _collect_matches(self, results: Dict[str, ProviderResult]) -> List[NormalizedMatch]:
        out: List[NormalizedMatch] = []
        for r in results.values():
            if not r.success or r.capability.value != "recognition":
                continue
            for m in r.normalized.get("matches", []):
                out.append(NormalizedMatch(
                    query_face_index=m.get("query_face_index", 0),
                    best_match=m.get("best_match"),
                    distance=m.get("distance", 1.0),
                    distances=m.get("distances", {}),
                    recognizer=r.provider,
                ))
        return out

    def _collect_scraped(self, results: Dict[str, ProviderResult]) -> List[NormalizedScrapeImage]:
        out: List[NormalizedScrapeImage] = []
        for r in results.values():
            if not r.success or r.capability.value != "scraping":
                continue
            for img in r.normalized.get("images", []):
                out.append(NormalizedScrapeImage(
                    url=img.get("url", ""),
                    alt=img.get("alt", ""),
                    source_page=img.get("source_page", ""),
                    scraper=r.provider,
                    width=img.get("width"),
                    height=img.get("height"),
                ))
        deduped = ResultPostprocessor.dedupe_images([s.__dict__ for s in out])
        return [NormalizedScrapeImage(**d) for d in deduped]

    def _collect_reverse(self, results: Dict[str, ProviderResult]) -> List[NormalizedReverseMatch]:
        out: List[NormalizedReverseMatch] = []
        for r in results.values():
            if not r.success or r.capability.value != "reverse_search":
                continue
            for m in r.normalized.get("results", []):
                out.append(NormalizedReverseMatch(
                    image_url=m.get("image_url", ""),
                    source_page=m.get("source_page", ""),
                    title=m.get("title", ""),
                    snippet=m.get("snippet", ""),
                    thumbnail=m.get("thumbnail", ""),
                    provider=r.provider,
                ))
        return out

    def _collect_image_analyses(self, results: Dict[str, ProviderResult]) -> List[NormalizedImageAnalysis]:
        out: List[NormalizedImageAnalysis] = []
        for r in results.values():
            if not r.success or r.capability.value != "image_analysis":
                continue
            n = r.normalized
            out.append(NormalizedImageAnalysis(
                provider=r.provider,
                quality_score=n.get("quality_score"),
                brightness=n.get("brightness"),
                contrast=n.get("contrast"),
                sharpness=n.get("sharpness"),
                noise_level=n.get("noise_level"),
                width=n.get("width"),
                height=n.get("height"),
                channels=n.get("channels"),
                color_profile=n.get("color_profile"),
                dominant_colors=n.get("dominant_colors", []),
                aspects=n.get("aspects", {}),
            ))
        return out

    def _collect_metadata(self, results: Dict[str, ProviderResult]) -> List[NormalizedMetadata]:
        out: List[NormalizedMetadata] = []
        for r in results.values():
            if not r.success or r.capability.value != "metadata":
                continue
            n = r.normalized
            out.append(NormalizedMetadata(
                provider=r.provider,
                format=n.get("format"),
                exif=n.get("exif", {}),
                xmp=n.get("xmp", {}),
                iptc=n.get("iptc", {}),
                gps=n.get("gps"),
                camera_make=n.get("camera_make"),
                camera_model=n.get("camera_model"),
                software=n.get("software"),
                capture_time=n.get("capture_time"),
            ))
        return out

    def _collect_forensics(self, results: Dict[str, ProviderResult]) -> List[NormalizedForensics]:
        out: List[NormalizedForensics] = []
        for r in results.values():
            if not r.success or r.capability.value != "forensics":
                continue
            n = r.normalized
            out.append(NormalizedForensics(
                provider=r.provider,
                integrity_score=n.get("integrity_score"),
                is_duplicate=n.get("is_duplicate"),
                duplicate_of=n.get("duplicate_of"),
                similarity_score=n.get("similarity_score"),
                manipulation_indicators=n.get("manipulation_indicators", []),
                ela_score=n.get("ela_score"),
                noise_inconsistency=n.get("noise_inconsistency"),
                details=n.get("details", {}),
            ))
        return out

    def _collect_ocr(self, results: Dict[str, ProviderResult]) -> List[OCRResult]:
        out: List[OCRResult] = []
        for r in results.values():
            if not r.success or r.capability.value != "ocr":
                continue
            n = r.normalized
            out.append(OCRResult(
                provider=r.provider,
                text_blocks=n.get("text_blocks", []),
                full_text=n.get("full_text", ""),
                language=n.get("language"),
            ))
        return out

    def _collect_object_detections(self, results: Dict[str, ProviderResult]) -> List[ObjectDetectionResult]:
        out: List[ObjectDetectionResult] = []
        for r in results.values():
            if not r.success or r.capability.value != "object_detection":
                continue
            n = r.normalized
            out.append(ObjectDetectionResult(
                provider=r.provider,
                objects=n.get("objects", []),
                model=n.get("model"),
            ))
        return out

    def _collect_scene(self, results: Dict[str, ProviderResult]) -> List[SceneResult]:
        out: List[SceneResult] = []
        for r in results.values():
            if not r.success or r.capability.value != "scene_recognition":
                continue
            n = r.normalized
            out.append(SceneResult(
                provider=r.provider,
                labels=n.get("labels", []),
            ))
        return out

    def _collect_nsfw(self, results: Dict[str, ProviderResult]) -> List[NSFWResult]:
        out: List[NSFWResult] = []
        for r in results.values():
            if not r.success or r.capability.value != "nsfw_detection":
                continue
            n = r.normalized
            out.append(NSFWResult(
                provider=r.provider,
                is_nsfw=n.get("is_nsfw", False),
                labels=n.get("labels", []),
            ))
        return out

    def _collect_ai_detection(self, results: Dict[str, ProviderResult]) -> List[AIDetectionResult]:
        out: List[AIDetectionResult] = []
        for r in results.values():
            if not r.success or r.capability.value != "ai_image_detection":
                continue
            n = r.normalized
            out.append(AIDetectionResult(
                provider=r.provider,
                is_ai_generated=n.get("is_ai_generated", False),
            ))
        return out

    def _collect_embeddings(self, results: Dict[str, ProviderResult]) -> List[EmbeddingResult]:
        out: List[EmbeddingResult] = []
        for r in results.values():
            if not r.success or r.capability.value != "embedding":
                continue
            n = r.normalized
            out.append(EmbeddingResult(
                provider=r.provider,
                embedding=n.get("embedding", []),
                model=n.get("model"),
                dimensions=n.get("dimensions", 0),
            ))
        return out

    def _build_evidence(self, results: Dict[str, ProviderResult]) -> List[Evidence]:
        out: List[Evidence] = []
        for r in results.values():
            out.append(Evidence(
                provider=r.provider,
                capability=r.capability.value,
                raw=r.raw,
                normalized=r.normalized,
                elapsed_ms=r.elapsed_ms,
                success=r.success,
                error=r.error,
                error_type=r.error_type,
                metadata=r.metadata,
                retry_count=r.retry_count,
            ))
        return out