"""
Tests for the resize-robust watermarking algorithm ('fm' mode).

The 'fm' algorithm uses DWT-QIM with delta=192, calibrated so that
the ±35% resize drift (std≈12.5 LL2 units) stays well within the
QIM tolerance of delta/4=48 units.

Covers:
  - Round-trip on clean PNG (0 errors required)
  - Round-trip after JPEG compression (0 errors required)
  - Round-trip after 80%/90%/120% resize (0 errors required)
  - No false-positive on unwatermarked images
  - Full metadata round-trip via encode/decode
  - High-level detect(algorithm='fm') and detect(algorithm='auto')
"""

import io
import numpy as np
import pytest
from PIL import Image

from clawid.algorithms.fourier_mellin import embed_bits, extract_bits, DELTA, REPEAT
from clawid.core.payload import encode, decode, to_bits, from_bits, HEADER_BYTES
from clawid.core.detect import detect
from clawid.core.embed import embed


# ── helpers ───────────────────────────────────────────────────────────────────

def _random_image(seed=42, size=(512, 512)) -> Image.Image:
    rng = np.random.default_rng(seed)
    arr = rng.integers(0, 256, (*size, 3), dtype=np.uint8)
    return Image.fromarray(arr, 'RGB')


def _jpeg_roundtrip(image: Image.Image, quality: int = 85) -> Image.Image:
    buf = io.BytesIO()
    image.save(buf, format='JPEG', quality=quality, subsampling=0)
    buf.seek(0)
    return Image.open(buf).convert('RGB')


def _resize_roundtrip(image: Image.Image, factor: float) -> Image.Image:
    orig = image.size
    new_w = int(orig[0] * factor)
    new_h = int(orig[1] * factor)
    resized = image.resize((new_w, new_h), Image.LANCZOS)
    return resized.resize(orig, Image.LANCZOS)


def _make_payload_bits() -> list:
    meta = {'uid': 'fm_test', 'platform': 'clawid-tests'}
    return to_bits(encode(meta))


# ── basic round-trip ──────────────────────────────────────────────────────────

def test_fm_embed_extract_png():
    """Embed → extract from PIL image (no save): 0 bit errors."""
    image = _random_image()
    bits = _make_payload_bits()

    wm = embed_bits(image, bits)
    extracted = extract_bits(wm, len(bits))

    errors = sum(a != b for a, b in zip(bits, extracted))
    assert errors == 0, f"{errors} bit errors on PNG round-trip"


def test_fm_embed_extract_jpeg_q85():
    """Embed → JPEG q=85 → extract: 0 bit errors (fm uses delta=192)."""
    image = _random_image(seed=7)
    bits = _make_payload_bits()

    wm = embed_bits(image, bits)
    wm_jpeg = _jpeg_roundtrip(wm, quality=85)
    extracted = extract_bits(wm_jpeg, len(bits))

    errors = sum(a != b for a, b in zip(bits, extracted))
    assert errors == 0, f"{errors} bit errors after JPEG q=85"


def test_fm_embed_extract_jpeg_q75():
    """Embed → JPEG q=75 → extract: low BER (moderate compression)."""
    image = _random_image(seed=13)
    bits = _make_payload_bits()

    wm = embed_bits(image, bits)
    wm_jpeg = _jpeg_roundtrip(wm, quality=75)
    extracted = extract_bits(wm_jpeg, len(bits))

    errors = sum(a != b for a, b in zip(bits, extracted))
    ber = errors / len(bits)
    assert ber < 0.05, f"BER {ber:.2%} after JPEG q=75 (expected <5%)"


# ── resize robustness ─────────────────────────────────────────────────────────

def test_fm_resize_80_percent():
    """After 80% resize → restore: 0 bit errors with fm (delta=192)."""
    image = _random_image(seed=21)
    bits = _make_payload_bits()

    wm = embed_bits(image, bits)
    wm_resized = _resize_roundtrip(wm, 0.80)
    extracted = extract_bits(wm_resized, len(bits))

    errors = sum(a != b for a, b in zip(bits, extracted))
    assert errors == 0, f"{errors} bit errors after 80% resize"


def test_fm_resize_90_percent():
    """After 90% resize → restore: 0 bit errors."""
    image = _random_image(seed=22)
    bits = _make_payload_bits()

    wm = embed_bits(image, bits)
    wm_resized = _resize_roundtrip(wm, 0.90)
    extracted = extract_bits(wm_resized, len(bits))

    errors = sum(a != b for a, b in zip(bits, extracted))
    assert errors == 0, f"{errors} bit errors after 90% resize"


def test_fm_resize_120_percent():
    """After 120% upscale → restore: 0 bit errors."""
    image = _random_image(seed=23)
    bits = _make_payload_bits()

    wm = embed_bits(image, bits)
    wm_resized = _resize_roundtrip(wm, 1.20)
    extracted = extract_bits(wm_resized, len(bits))

    errors = sum(a != b for a, b in zip(bits, extracted))
    assert errors == 0, f"{errors} bit errors after 120% resize"


# ── no false positives ────────────────────────────────────────────────────────

def test_fm_no_false_positive():
    """Unwatermarked image should not decode a valid MAGIC header."""
    image = _random_image(seed=99)
    header_bits = extract_bits(image, HEADER_BYTES * 8)
    header_bytes = from_bits(header_bits)

    from clawid.core.payload import MAGIC, VERSION
    assert header_bytes[:4] != MAGIC or header_bytes[4] != VERSION, \
        "False positive: random image decoded a valid clawID header"


# ── full metadata round-trip ──────────────────────────────────────────────────

def test_fm_full_metadata_roundtrip():
    """embed_bits → extract_bits → decode: metadata survives 80% resize."""
    image = _random_image(seed=55)
    meta = {'uid': 'alice', 'platform': 'hawky.ai', 'asset_id': 'img-001'}
    payload = encode(meta)
    bits = to_bits(payload)

    wm = embed_bits(image, bits)
    wm_r = _resize_roundtrip(wm, 0.80)
    extracted_bits = extract_bits(wm_r, len(bits))
    recovered = decode(from_bits(extracted_bits))

    assert recovered['uid'] == meta['uid']
    assert recovered['platform'] == meta['platform']
    assert recovered['asset_id'] == meta['asset_id']


# ── high-level API ────────────────────────────────────────────────────────────

def test_detect_fm_algorithm(tmp_path):
    """High-level embed(algorithm='fm') + detect(algorithm='fm') works."""
    src = _random_image(seed=77)
    out = str(tmp_path / "wm_fm.png")

    meta = embed(src, out, {'uid': 'bob', 'platform': 'hawky.ai'}, algorithm='fm')
    result = detect(out, algorithm='fm')

    assert result is not None, "detect() returned None for FM-watermarked image"
    assert result['uid'] == 'bob'
    assert result['clawid'] == meta['clawid']


def test_detect_auto_falls_back_to_fm(tmp_path):
    """Auto detect finds FM-embedded watermark (QIM check fails, FM succeeds)."""
    src = _random_image(seed=88)
    out = str(tmp_path / "wm_fm_auto.png")

    meta = embed(src, out, {'uid': 'carol'}, algorithm='fm')
    result = detect(out, algorithm='auto')

    assert result is not None, "auto detect() failed to find FM watermark"
    assert result['uid'] == 'carol'


def test_detect_fm_after_resize(tmp_path):
    """FM-embedded image detected correctly after 80% resize."""
    src = _random_image(seed=66)
    out = str(tmp_path / "wm_fm_resize.png")

    meta = embed(src, out, {'uid': 'dave'}, algorithm='fm')

    wm = Image.open(out)
    wm_r = _resize_roundtrip(wm, 0.80)
    wm_r_path = str(tmp_path / "wm_fm_resized.png")
    wm_r.save(wm_r_path)

    result = detect(wm_r_path, algorithm='fm')
    assert result is not None, "detect() failed after 80% resize"
    assert result['uid'] == 'dave'
