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

from clawid.core.embed import embed
from clawid.core.detect import detect
from clawid.algorithms.dwt_qim import embed_bits, MIN_DIM


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

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


def _embed_detect(tmp_path, metadata, ext='png', image_size=(512, 512)):
    img = make_image(image_size)
    src = str(tmp_path / f'src.{ext}')
    out = str(tmp_path / f'out.{ext}')
    img.save(src)
    embed(src, out, metadata)
    return detect(out)


# ── tests ─────────────────────────────────────────────────────────────────────

def test_embed_detect_png(tmp_path):
    meta = {'uid': 'user123', 'platform': 'hawky.ai'}
    result = _embed_detect(tmp_path, meta, ext='png')
    assert result is not None
    assert result['uid'] == 'user123'
    assert result['platform'] == 'hawky.ai'
    assert 'clawid' in result
    assert 'ts' in result


def test_embed_detect_jpeg(tmp_path):
    """Watermark should survive JPEG compression at quality=95."""
    meta = {'uid': 'jpeguser', 'platform': 'test'}
    result = _embed_detect(tmp_path, meta, ext='jpg')
    assert result is not None
    assert result['uid'] == 'jpeguser'


def test_no_watermark_returns_none(tmp_path):
    img = make_image()
    path = str(tmp_path / 'clean.png')
    img.save(path)
    assert detect(path) is None


def test_custom_clawid_preserved(tmp_path):
    meta = {'uid': 'u1', 'clawid': 'my-custom-id-123'}
    result = _embed_detect(tmp_path, meta)
    assert result['clawid'] == 'my-custom-id-123'


def test_auto_clawid_generated(tmp_path):
    meta = {'uid': 'u2'}
    result = _embed_detect(tmp_path, meta)
    assert 'clawid' in result
    assert len(result['clawid']) == 36  # UUID format


def test_image_not_modified_visually(tmp_path):
    """Invisible watermark should not visibly change image (PSNR > 40 dB)."""
    img = make_image((512, 512))
    src = str(tmp_path / 'src.png')
    out = str(tmp_path / 'out.png')
    img.save(src)
    embed(src, out, {'uid': 'psnr_test'})

    orig = np.array(img, dtype=np.float64)
    watermarked = np.array(Image.open(out), dtype=np.float64)
    mse = np.mean((orig - watermarked) ** 2)
    if mse == 0:
        return  # identical — even better
    psnr = 10 * np.log10(255 ** 2 / mse)
    assert psnr > 40.0, f"PSNR too low: {psnr:.1f} dB"


def test_too_small_image_raises():
    tiny = Image.new('RGB', (MIN_DIM - 1, MIN_DIM - 1))
    with pytest.raises(ValueError, match=str(MIN_DIM)):
        embed_bits(tiny, [0, 1])


def test_detect_from_pil_image(tmp_path):
    """detect() should accept a PIL Image directly, not just a file path."""
    meta = {'uid': 'pil_test'}
    src = str(tmp_path / 'src.png')
    out = str(tmp_path / 'out.png')
    make_image().save(src)
    embed(src, out, meta)
    pil_img = Image.open(out)
    result = detect(pil_img)
    assert result is not None
    assert result['uid'] == 'pil_test'
