"""
QR-Mark: AI-readable invisible watermarking via tiled semi-transparent QR codes.

Embeds metadata as a 2×2 tiled QR code at ~15% opacity.

Properties:
  - Invisible to casual human viewing at default opacity
  - Directly readable by AI vision models (GPT-4V, Gemini, Grok, Claude)
    — just ask: "what QR code or watermark is in this image?"
  - Readable by any smartphone QR scanner app
  - Survives: screenshot, JPEG q≥70, up to 50% crop (2×2 tiling)
  - Works on images AND video frames

Tiling strategy: 2×2 grid, each tile = 50% image dimensions.
Even if the image is cropped to any single quadrant, one full QR remains.
"""

import json
import numpy as np
from PIL import Image, ImageEnhance, ImageFilter
from typing import Optional

try:
    import qrcode
    import qrcode.constants
    _QRCODE_OK = True
except ImportError:
    _QRCODE_OK = False

try:
    import cv2
    _CV2_OK = True
except ImportError:
    _CV2_OK = False

try:
    from pyzbar import pyzbar as _pyzbar
    _PYZBAR_OK = True
except ImportError:
    _PYZBAR_OK = False


DEFAULT_OPACITY = 0.18   # 18% — imperceptible to humans, reliable for AI vision
MIN_QR_DIM = 100         # minimum pixel size per QR tile to stay decodable


def embed_qr(
    image: Image.Image,
    metadata: dict,
    opacity: float = DEFAULT_OPACITY,
) -> Image.Image:
    """
    Embed metadata as a tiled semi-transparent QR code.

    Args:
        image:    RGB PIL Image (min 256×256)
        metadata: dict — will be serialised as compact JSON into the QR
        opacity:  0.0–1.0 blend strength (default 0.15)
                  0.10 = very subtle, AI needs good preprocessing
                  0.15 = default, reliable for AI vision models
                  0.25 = clearly visible, maximum scanner compatibility

    Returns:
        Watermarked RGB PIL Image
    """
    if not _QRCODE_OK:
        raise ImportError(
            "qrcode package required: pip install 'qrcode[pil]'"
        )

    img = image.convert('RGB')
    w, h = img.size

    # Compact QR payload — short keys reduce QR version → larger modules → better detection
    _KEY_MAP = {'clawid': 'c', 'uid': 'u', 'platform': 'p', 'asset_id': 'a', 'ts': 't'}
    compact = {_KEY_MAP.get(k, k): v for k, v in metadata.items()}
    payload = json.dumps(compact, separators=(',', ':'), ensure_ascii=False)

    # Generate QR with HIGH error correction (survives 30% module damage)
    qr = qrcode.QRCode(
        version=None,                              # auto-pick smallest version
        error_correction=qrcode.constants.ERROR_CORRECT_H,
        box_size=10,
        border=3,
    )
    qr.add_data(payload)
    qr.make(fit=True)
    qr_raw = qr.make_image(fill_color='black', back_color='white').convert('L')

    # Each tile covers one quadrant (w//2 × h//2)
    tile_w = max(w // 2, MIN_QR_DIM)
    tile_h = max(h // 2, MIN_QR_DIM)

    # Resize QR to tile size — LANCZOS then threshold for clean modules
    qr_tile = qr_raw.resize((tile_w, tile_h), Image.LANCZOS)
    # Re-binarise after resize to keep crisp black/white modules
    qr_tile = qr_tile.point(lambda p: 0 if p < 128 else 255, '1').convert('L')

    # Build 2×2 tiled overlay at full image resolution
    tiled = Image.new('L', (w, h), 255)
    for row in range(2):
        for col in range(2):
            tiled.paste(qr_tile, (col * tile_w, row * tile_h))

    # Linear blend: QR black modules darken, white modules lighten
    # This preserves BOTH dark and light QR modules so decoders can read it
    # result = orig * (1 - opacity) + qr_bw * opacity
    orig_arr  = np.array(img, dtype=np.float32)               # H×W×3
    tiled_rgb = np.stack([np.array(tiled)] * 3, axis=2).astype(np.float32)  # H×W×3
    blended   = orig_arr * (1.0 - opacity) + tiled_rgb * opacity
    result    = Image.fromarray(np.clip(blended, 0, 255).astype(np.uint8), 'RGB')
    return result


def detect_qr(image: Image.Image) -> Optional[dict]:
    """
    Detect and decode the QR watermark from an image.

    Tries multiple preprocessing strategies to maximise detection rate
    even after JPEG compression, screenshots, or colour shifts.

    Returns:
        Metadata dict if a valid clawID QR watermark is found, None otherwise.
    """
    if not (_CV2_OK or _PYZBAR_OK):
        raise ImportError(
            "OpenCV or pyzbar required: "
            "pip install opencv-python-headless   OR   pip install pyzbar"
        )

    arr = np.array(image.convert('RGB'))

    # Build a pipeline of increasingly aggressive preprocessing variants
    # Most watermarked images will decode on the first or second pass
    variants = _preprocessing_pipeline(arr)

    for variant in variants:
        result = _try_decode(variant)
        if result is not None:
            return result

    return None


# ── internal helpers ──────────────────────────────────────────────────────────

def _try_decode(arr: np.ndarray) -> Optional[dict]:
    """Attempt QR decode with pyzbar (preferred) then OpenCV."""
    gray = _to_gray_cv(arr) if _CV2_OK else _to_gray_np(arr)

    # ── pyzbar ────────────────────────────────────────────────────────────────
    if _PYZBAR_OK:
        codes = _pyzbar.decode(Image.fromarray(gray))
        for code in codes:
            if code.type == 'QRCODE':
                data = _parse_payload(code.data)
                if data is not None:
                    return data

    # ── OpenCV ────────────────────────────────────────────────────────────────
    if _CV2_OK:
        detector = cv2.QRCodeDetector()
        # Standard detect
        raw, _, _ = detector.detectAndDecode(gray)
        if raw:
            data = _parse_payload(raw.encode())
            if data is not None:
                return data
        # detectAndDecodeMulti catches tiled copies
        ok, raws, _, _ = detector.detectAndDecodeMulti(gray)
        if ok:
            for raw in raws:
                if raw:
                    data = _parse_payload(raw.encode())
                    if data is not None:
                        return data

    return None


_KEY_EXPAND = {'c': 'clawid', 'u': 'uid', 'p': 'platform', 'a': 'asset_id', 't': 'ts'}


def _parse_payload(raw: bytes) -> Optional[dict]:
    """Parse QR payload bytes → metadata dict (expands compact keys)."""
    try:
        text = raw.decode('utf-8') if isinstance(raw, bytes) else raw
        data = json.loads(text)
        # Expand compact keys back to full names
        return {_KEY_EXPAND.get(k, k): v for k, v in data.items()}
    except (json.JSONDecodeError, UnicodeDecodeError, ValueError):
        return None


def _preprocessing_pipeline(arr: np.ndarray):
    """
    Generate preprocessing variants from least to most aggressive.
    Each is a numpy uint8 RGB array.

    Key insight: we tile the QR 2×2, so averaging the 4 quadrants
    reinforces the QR signal (same in all 4) while cancelling the
    background image (different in each quadrant). This dramatically
    improves detection after JPEG compression or on noisy images.
    """
    variants = [arr]  # 1. raw

    # 2. Tile-averaged: average all 4 quadrants → QR signal coherently adds,
    #    background content partially cancels → best for JPEG/noisy images
    h, w = arr.shape[:2]
    th, tw = h // 2, w // 2
    if th >= MIN_QR_DIM and tw >= MIN_QR_DIM:
        q1 = arr[:th,   :tw  ].astype(np.float32)
        q2 = arr[:th,   tw:tw*2].astype(np.float32) if tw*2 <= w else q1
        q3 = arr[th:th*2, :tw  ].astype(np.float32) if th*2 <= h else q1
        q4 = arr[th:th*2, tw:tw*2].astype(np.float32) if th*2 <= h and tw*2 <= w else q1
        averaged = np.clip((q1 + q2 + q3 + q4) / 4, 0, 255).astype(np.uint8)
        variants.append(averaged)

        # 3. Tile-averaged + CLAHE contrast enhancement
        if _CV2_OK:
            lab = cv2.cvtColor(averaged, cv2.COLOR_RGB2LAB)
            lc, a, b = cv2.split(lab)
            clahe = cv2.createCLAHE(clipLimit=4.0, tileGridSize=(4, 4))
            enhanced = cv2.merge([clahe.apply(lc), a, b])
            variants.append(cv2.cvtColor(enhanced, cv2.COLOR_LAB2RGB))

            # 4. Tile-averaged + median blur (removes texture noise, preserves QR modules)
            blurred = cv2.medianBlur(averaged, 5)
            variants.append(blurred)

            # 5. Tile-averaged + median blur + CLAHE
            lab2 = cv2.cvtColor(blurred, cv2.COLOR_RGB2LAB)
            lc2, a2, b2 = cv2.split(lab2)
            enhanced2 = cv2.merge([clahe.apply(lc2), a2, b2])
            variants.append(cv2.cvtColor(enhanced2, cv2.COLOR_LAB2RGB))

            # 6. Tile-averaged + Otsu threshold
            gray_avg = _to_gray_cv(blurred)
            _, otsu = cv2.threshold(gray_avg, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
            variants.append(cv2.cvtColor(otsu, cv2.COLOR_GRAY2RGB))

            # 7. Tile-averaged + upscale + sharpen
            up = cv2.resize(averaged, (tw * 4, th * 4), interpolation=cv2.INTER_CUBIC)
            sharp_k = np.array([[-1,-1,-1],[-1,9,-1],[-1,-1,-1]], dtype=np.float32)
            variants.append(cv2.filter2D(up, -1, sharp_k))

    if not _CV2_OK:
        return variants

    # 6. Original + sharpened (fallback if tiling doesn't help)
    sharp_k = np.array([[-1,-1,-1],[-1,9,-1],[-1,-1,-1]], dtype=np.float32)
    variants.append(cv2.filter2D(arr, -1, sharp_k))

    # 7. Original + CLAHE
    lab = cv2.cvtColor(arr, cv2.COLOR_RGB2LAB)
    lc, a, b = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
    enhanced = cv2.merge([clahe.apply(lc), a, b])
    variants.append(cv2.cvtColor(enhanced, cv2.COLOR_LAB2RGB))

    return variants


def _to_gray_cv(arr: np.ndarray) -> np.ndarray:
    return cv2.cvtColor(arr, cv2.COLOR_RGB2GRAY)


def _to_gray_np(arr: np.ndarray) -> np.ndarray:
    return (0.299 * arr[:, :, 0] + 0.587 * arr[:, :, 1]
            + 0.114 * arr[:, :, 2]).astype(np.uint8)
