"""
NyxAI FORENSIC DETECTOR - HARDWARE & MODEL FINGERPRINTING ENGINE
File: engine.py
Features: Byte-Level Header Inspection, FFT Spectral Kurtosis, VAE Noise Profiling
"""

import os
import io
import json
import re
from typing import List, Dict, Any, Tuple
import numpy as np
import cv2
from PIL import Image, ImageDraw, ImageChops, ImageEnhance, ExifTags
from google import genai
from google.genai import types

# =============================================================================
# 1. ADVANCED MATHEMATICAL & BYTE-LEVEL FINGERPRINT EXTRACTION
# =============================================================================

class AdvancedFingerprintExtractor:
    @staticmethod
    def inspect_raw_bytes_and_exif(image: Image.Image) -> Dict[str, Any]:
        """Scans image byte stream and EXIF for digital fingerprints (Adobe, C2PA, SynthID, etc.)"""
        fingerprints = []
        software_detected = "None"
        
        try:
            buf = io.BytesIO()
            image.save(buf, format=image.format or "JPEG")
            raw_bytes = buf.getvalue().lower()

            # Known Byte Signatures
            if b"synthid" in raw_bytes or b"google" in raw_bytes:
                fingerprints.append("Google SynthID Watermark Signature Detected")
                software_detected = "Google Imagen / Gemini"
            if b"dall-e" in raw_bytes or b"openai" in raw_bytes:
                fingerprints.append("OpenAI / DALL-E Metadata Header Marker Detected")
                software_detected = "OpenAI DALL-E 3"
            if b"midjourney" in raw_bytes:
                fingerprints.append("Midjourney Software Header Tag Detected")
                software_detected = "Midjourney"
            if b"adobe" in raw_bytes or b"photoshop" in raw_bytes:
                fingerprints.append("Adobe Generative Fill / Firefly Tag Detected")
                software_detected = "Adobe Firefly / Photoshop AI"
            if b"comfyui" in raw_bytes or b"automatic1111" in raw_bytes or b"stablediffusion" in raw_bytes:
                fingerprints.append("Stable Diffusion / ComfyUI VAE Pipeline Signature Detected")
                software_detected = "Stable Diffusion XL / 1.5"

        except Exception:
            pass

        return {
            "detected_raw_software": software_detected,
            "byte_fingerprint_markers": fingerprints
        }

    @staticmethod
    def compute_advanced_signal_metrics(image: Image.Image) -> Tuple[Dict[str, Image.Image], Dict[str, Any], Dict[str, str]]:
        rgb = image.convert('RGB')
        arr = np.array(rgb)
        gray = cv2.cvtColor(arr, cv2.COLOR_RGB2GRAY)
        h, w = gray.shape

        # 1. ELA Computation & Compression Delta
        buf = io.BytesIO()
        rgb.save(buf, 'JPEG', quality=90)
        buf.seek(0)
        resaved = Image.open(buf)
        diff = ImageChops.difference(rgb, resaved)
        diff_arr = np.array(diff, dtype=np.float32)
        ela_mean = float(np.mean(diff_arr))
        ela_std = float(np.std(diff_arr))
        ela_score = round(float(min(100.0, (ela_mean / 10.0) * 100.0)), 2)

        extrema = diff.getextrema()
        max_diff = max([ex[1] for ex in extrema]) or 1
        ela_img = ImageEnhance.Brightness(diff).enhance((255.0 / max_diff) * 1.5)

        # 2. 2D FFT Kurtosis & High-Frequency Anomaly Index (Pure NumPy)
        f = np.fft.fft2(gray)
        fshift = np.fft.fftshift(f)
        magnitude = 20 * np.log(np.abs(fshift) + 1e-5)
        
        center_h, center_w = h // 2, w // 2
        mag_hp = magnitude.copy()
        cv2.circle(mag_hp, (center_w, center_h), min(h, w) // 8, 0, -1)
        
        flat_hp = mag_hp.ravel()
        fft_std = float(np.std(flat_hp)) or 1e-5
        fft_mean = float(np.mean(flat_hp))
        # Pure NumPy Kurtosis calculation
        fft_kurt = float(np.mean(((flat_hp - fft_mean) / fft_std) ** 4) - 3.0)
        fft_anomaly_score = round(float(min(100.0, (fft_std / 12.0) * 100.0)), 2)

        norm_fft = cv2.normalize(magnitude, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
        fft_colored = cv2.applyColorMap(norm_fft, cv2.COLORMAP_TURBO)
        fft_img = Image.fromarray(cv2.cvtColor(fft_colored, cv2.COLOR_BGR2RGB))

        # 3. Laplacian Sharpness & Edge Distribution
        lap = cv2.Laplacian(gray, cv2.CV_8U, ksize=3)
        lap_var = round(float(cv2.Laplacian(gray, cv2.CV_64F).var()), 2)
        lap_img = Image.fromarray(cv2.cvtColor(cv2.applyColorMap(lap, cv2.COLORMAP_MAGMA), cv2.COLOR_BGR2RGB))

        # 4. Luminance Flow
        grad_x = cv2.Sobel(gray, cv2.CV_32F, 1, 0, ksize=3)
        grad_y = cv2.Sobel(gray, cv2.CV_32F, 0, 1, ksize=3)
        mag, _ = cv2.cartToPolar(grad_x, grad_y, angleInDegrees=True)
        lum_img = Image.fromarray(cv2.cvtColor(cv2.applyColorMap(cv2.normalize(mag, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U), cv2.COLORMAP_VIRIDIS), cv2.COLOR_BGR2RGB))

        # 5. RGB Disparity
        r, g, b = arr[:,:,0], arr[:,:,1], arr[:,:,2]
        disp = np.abs(r.astype(int) - g.astype(int)) + np.abs(g.astype(int) - b.astype(int))
        rgb_disp_img = Image.fromarray(cv2.cvtColor(cv2.applyColorMap(cv2.normalize(disp.astype(np.uint8), None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U), cv2.COLORMAP_JET), cv2.COLOR_BGR2RGB))

        # 6. Canny Boundary
        canny = cv2.Canny(gray, 80, 180)
        canny_img = Image.fromarray(cv2.cvtColor(cv2.applyColorMap(canny, cv2.COLORMAP_HOT), cv2.COLOR_BGR2RGB))

        # 7. Local Texture Variance
        blur = cv2.GaussianBlur(gray, (5, 5), 0)
        texture_diff = cv2.absdiff(gray, blur)
        tex_img = Image.fromarray(cv2.cvtColor(cv2.applyColorMap(cv2.normalize(texture_diff, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U), cv2.COLORMAP_CIVIDIS), cv2.COLOR_BGR2RGB))

        # 8. HSV Saturation
        hsv = cv2.cvtColor(arr, cv2.COLOR_RGB2HSV)
        sat_img = Image.fromarray(cv2.cvtColor(cv2.applyColorMap(hsv[:,:,1], cv2.COLORMAP_PLASMA), cv2.COLOR_BGR2RGB))

        # 9. High-Pass PRNU Sensor Noise
        denoised = cv2.fastNlMeansDenoising(gray, None, 10, 7, 21)
        noise_residue = cv2.absdiff(gray, denoised)
        prnu_img = Image.fromarray(cv2.cvtColor(cv2.applyColorMap(cv2.normalize(noise_residue, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U), cv2.COLORMAP_TWILIGHT), cv2.COLOR_BGR2RGB))

        maps = {
            "ela": ela_img, "fft": fft_img, "laplacian": lap_img, "luminance": lum_img,
            "rgb_disparity": rgb_disp_img, "canny": canny_img, "texture": tex_img,
            "hsv_sat": sat_img, "prnu_noise": prnu_img
        }

        raw_byte_data = AdvancedFingerprintExtractor.inspect_raw_bytes_and_exif(image)

        math_summary = {
            "fft_anomaly_score": f"{fft_anomaly_score}%",
            "fft_kurtosis_index": round(fft_kurt, 3),
            "ela_compression_variance": f"{ela_score}%",
            "ela_pixel_std_dev": round(ela_std, 2),
            "laplacian_sharpness_variance": lap_var,
            "raw_byte_software_tag": raw_byte_data["detected_raw_software"],
            "raw_byte_signatures": raw_byte_data["byte_fingerprint_markers"]
        }

        descriptions = {
            "spatial": "Identifies physical defects, bad typography, or anatomical errors with numbered bounding boxes.",
            "ela": f"Highlights digital re-compression stress. ELA Variance: {ela_score}%.",
            "fft": f"Scans high-frequency pixel distribution. FFT Anomaly: {fft_anomaly_score}%, Kurtosis: {round(fft_kurt, 2)}.",
            "laplacian": f"Evaluates edge sharpness across focal planes. Laplacian Var: {lap_var}.",
            "luminance": "Traces light particle angle and shadow vectors to detect artificial object insertion.",
            "rgb_disparity": "Measures color channel cross-talk variance.",
            "canny": "Exposes object boundary continuity and alpha matting gaps.",
            "texture": "Analyzes micro-surface grain consistency and generative smoothing.",
            "hsv_sat": "Maps color purity and unnatural saturation clustering.",
            "prnu": "Extracts high-frequency sensor noise to verify hardware camera fingerprints."
        }

        return maps, math_summary, descriptions

# =============================================================================
# 2. PROMPT FOR MULTI-MODEL PROBABILITY & FINGERPRINT ATTRIBUTION
# =============================================================================

STAGE_3_STRICT_PROMPT = """
You are the Lead Generative AI Forensics Auditor at NyxAI.

PERCENTAGE-BASED SCORING DEFINITION:
- final_weighted_ai_score: 0 to 100% (where 100% = 100% Confirmed Synthetic/AI Generated, and 0% = 100% Genuine Camera Photograph).

MODEL FINGERPRINT RECOGNITION TAXONOMY (Evaluate probability share for each):
1. Google Imagen / Gemini Flash Image (Nano Banana): Look for ultra-smooth lighting, realistic skin pores, SynthID frequency footprints, natural depth-of-field.
2. OpenAI DALL-E 3 / ChatGPT Vision: Look for hyper-saturated colors, perfect illustration rendering, vector-like text, smooth waxy surfaces.
3. Midjourney (v5 / v6 / v6.1): Look for artistic film grain, painterly micro-textures, hyper-detailed eyes and hair strand anti-aliasing, cinema lighting.
4. Black Forest Labs Flux.1 (Dev / Pro / Schnell): Look for flawless photorealistic text rendering, extremely accurate hands/teeth without waxy sheen.
5. Stability AI Stable Diffusion (SDXL / SD 3.5): Look for distinct high-frequency grid noise in dark areas, character alignment inconsistencies, VAE color banding.
6. Adobe Firefly / Photoshop Generative Fill: Look for localized edge-matting seams, background lighting mismatch around edited objects.
7. Alibaba Qwen / Wan 2.1 / Kling / Luma: Video/Hybrid frame interpolation patterns.

MANDATED REQUIREMENT:
Output a structured list of candidates in `model_attributions` with their probability shares (must sum up to 100% if AI generated, or indicate 0% if real).

Return ONLY raw JSON with NO markdown code blocks:
{
  "final_weighted_ai_score": <number 0-100>,
  "confidence_interval": "±<1-5>%",
  "final_verdict": "<CONFIRMED_AI_GENERATED / CONFIRMED_REAL_PHOTOGRAPH / HIGHLY_SUSPICIOUS>",
  "is_hybrid_multi_model": <true / false>,
  "model_attributions": [
    {
      "model_name": "<Exact Model Name e.g. Midjourney v6.1 / Google Gemini / OpenAI DALL-E 3 / Flux.1 / Stable Diffusion XL / Adobe Firefly>",
      "probability_share": <number 0-100>,
      "fingerprint_evidence": "<Specific visual or mathematical fingerprint proving this candidate>"
    }
  ],
  "domain_scores": {
    "frequency_noise": {"score": <0-100>, "proof": "<scientific evidence>"},
    "photometric_lighting": {"score": <0-100>, "proof": "<scientific evidence>"},
    "biological_anatomy": {"score": <0-100>, "proof": "<scientific evidence>"},
    "textural_homogeneity": {"score": <0-100>, "proof": "<scientific evidence>"},
    "optical_lens_physics": {"score": <0-100>, "proof": "<scientific evidence>"},
    "edge_matting": {"score": <0-100>, "proof": "<scientific evidence>"},
    "semantic_logic": {"score": <0-100>, "proof": "<scientific evidence>"},
    "typography_glyphs": {"score": <0-100>, "proof": "<scientific evidence>"},
    "color_space": {"score": <0-100>, "proof": "<scientific evidence>"},
    "generative_signature": {"score": <0-100>, "proof": "<scientific evidence>"}
  },
  "spatial_anomalies": [
    {
      "id": 1,
      "severity": "<CRITICAL / HIGH / MEDIUM>",
      "bounding_box": [ymin, xmin, ymax, xmax],
      "domain": "<Domain Name>",
      "description": "<Clear explanation>"
    }
  ],
  "executive_summary": "<Comprehensive technical summary justifying the probability score and model fingerprint share>"
}
"""

class SpatialAnnotator:
    @staticmethod
    def draw(image: Image.Image, anomalies: List[Dict[str, Any]]) -> Image.Image:
        annotated = image.copy().convert("RGB")
        draw = ImageDraw.Draw(annotated)
        w, h = annotated.size

        for item in anomalies:
            aid = item.get("id", 1)
            box = item.get("bounding_box", [0, 0, 0, 0])
            sev = item.get("severity", "HIGH")
            ymin, xmin, ymax, xmax = box

            color = "#E879F9" if sev == "CRITICAL" else ("#C084FC" if sev == "HIGH" else "#38BDF8")

            left, top = (xmin / 1000.0) * w, (ymin / 1000.0) * h
            right, bottom = (xmax / 1000.0) * w, (ymax / 1000.0) * h

            draw.rectangle([left, top, right, bottom], outline=color, width=4)
            r = 15
            draw.ellipse([left - r, top - r, left + r, top + r], fill=color)
            draw.text((left - 5, top - 8), str(aid), fill="#FFFFFF")

        return annotated

class ForensicEngine:
    def __init__(self):
        api_key = os.environ.get("GEMINI_API_KEY", "")
        if not api_key:
            raise ValueError("GEMINI_API_KEY environment variable is missing!")
        self.client = genai.Client(api_key=api_key)

    def process(self, image: Image.Image) -> Tuple[Image.Image, Dict[str, Image.Image], Dict[str, str], Dict[str, Any]]:
        # 1. Extract Signal Maps & Math Metrics
        signal_maps, math_summary, map_descs = AdvancedFingerprintExtractor.compute_advanced_signal_metrics(image)

        # 2. Feed Math + Byte Signatures to Gemini
        context_prompt = f"""
        EXPLICIT HARDWARE & SIGNAL MATHEMATICAL FINGERPRINTS EXTRACTED BY PYTHON:
        {json.dumps(math_summary, indent=2)}

        Execute a 100+ Model Fingerprint Recognition and 10-Domain Forensic Audit. Use the extracted mathematical noise index and raw byte markers as absolute ground truth.
        """

        audit_raw = self.client.models.generate_content(
            model="gemini-2.5-flash",
            contents=[image, context_prompt],
            config=types.GenerateContentConfig(system_instruction=STAGE_3_STRICT_PROMPT)
        )
        
        clean_text = audit_raw.text.strip()
        clean_text = re.sub(r'^```json\s*', '', clean_text)
        clean_text = re.sub(r'^```\s*', '', clean_text)
        clean_text = re.sub(r'\s*```$', '', clean_text)
        
        audit_data = json.loads(clean_text)

        spatial_map = SpatialAnnotator.draw(image, audit_data.get("spatial_anomalies", []))

        return spatial_map, signal_maps, map_descs, audit_data