"""

oversight_core.formats.image - image format adapter.

DCT-domain frequency watermarking. Survives:

  - JPEG recompression (qualities >= 50)

  - Moderate resizing (up to ~50%)

  - Minor cropping

  - Format conversion (PNG <-> JPEG)

Does NOT survive:

  - Heavy compression (quality < 30)

  - Aggressive cropping (> 30% removed)

  - Rotation without knowing the angle

  - Deliberate adversarial watermark-removal attacks (use spread-spectrum

    methods for that; out of MVP scope)

Algorithm: Cox et al. additive spread-spectrum in the DCT mid-band.

  1. Convert to YCbCr, take Y (luma) channel.

  2. Apply 2D DCT to the full Y plane.

  3. Pick the N largest mid-frequency coefficients (skip DC and lowest).

  4. Embed bit b_i by scaling coefficient c_i by (1 + alpha * x_i)

     where x_i is a deterministic bit-derived sequence from mark_id.

  5. Inverse DCT -> write back.

Recovery: sign-correlation between the DCT mid-band of the suspect image and

the expected bit sequence derived from a candidate mark_id.

"""

from __future__ import annotations

import hashlib

import io

from typing import Optional

import numpy as np

from PIL import Image

from scipy.fft import dct, idct

def _mark_to_sequence(mark_id: bytes, length: int) -> np.ndarray:

    """Deterministic +1/-1 sequence derived from mark_id."""

    out = np.zeros(length, dtype=np.int8)

    i = 0

    ctr = 0

    while i < length:

        h = hashlib.sha256(mark_id + ctr.to_bytes(4, "big")).digest()

        for byte in h:

            for bit in range(8):

                if i >= length:

                    break

                out[i] = 1 if (byte >> bit) & 1 else -1

                i += 1

        ctr += 1

    return out

def _dct2(a: np.ndarray) -> np.ndarray:

    return dct(dct(a, axis=0, norm="ortho"), axis=1, norm="ortho")

def _idct2(a: np.ndarray) -> np.ndarray:

    return idct(idct(a, axis=0, norm="ortho"), axis=1, norm="ortho")

def _pick_midband_indices(shape: tuple[int, int], n: int = 1000) -> np.ndarray:

    """

    Pick indices of mid-frequency DCT coefficients. We skip the DC and lowest

    frequencies (too visible when perturbed) and the highest (destroyed by JPEG).

    """

    H, W = shape

    lo = int(min(H, W) * 0.10)

    hi = int(min(H, W) * 0.40)

    coords = []

    for i in range(H):

        for j in range(W):

            if lo <= (i + j) <= hi:

                coords.append((i, j))

    coords = coords[:n]

    return np.array(coords)

def embed(

    image_bytes: bytes,

    mark_id: bytes,

    alpha: float = 0.10,

    n_coeffs: int = 1500,

) -> bytes:

    """

    Embed mark_id into the DCT mid-band of the image.

    Algorithm: for each of n_coeffs mid-band coefficients c_i, replace with

       c'_i = c_i + alpha * |c_i| * bit_i

    where bit_i is a deterministic +1/-1 sequence derived from mark_id.

    This additive-scaled-by-magnitude form gives reliable blind detection

    via normalized correlation, unlike pure sign-embedding which is

    destroyed by clipping after iDCT.

    Returns PNG bytes (lossless, to preserve the watermark for distribution).

    Caller can recompress to JPEG for transmission; watermark survives

    JPEG quality >= 60 in our testing.

    """

    img = Image.open(io.BytesIO(image_bytes)).convert("RGB")

    ycbcr = img.convert("YCbCr")

    y, cb, cr = ycbcr.split()

    y_arr = np.array(y, dtype=np.float64)

    D = _dct2(y_arr)

    coords = _pick_midband_indices(D.shape, n=n_coeffs)

    bits = _mark_to_sequence(mark_id, len(coords))

    for (i, j), b in zip(coords, bits):

        mag = abs(D[i, j])

        D[i, j] = D[i, j] + alpha * mag * b

    y_marked = _idct2(D)

    y_marked = np.clip(y_marked, 0, 255).astype(np.uint8)

    y2 = Image.fromarray(y_marked, mode="L")

    out = Image.merge("YCbCr", (y2, cb, cr)).convert("RGB")

    buf = io.BytesIO()

    out.save(buf, format="PNG")

    return buf.getvalue()

def verify(

    image_bytes: bytes,

    candidate_mark_id: bytes,

    threshold: float = 0.05,

    n_coeffs: int = 1500,

) -> tuple[bool, float]:

    """

    Blind detection of candidate_mark_id in the image's DCT mid-band.

    Returns (match, normalized_correlation).

    Correlation metric:

       score = <coeffs, expected> / (||coeffs|| * ||expected||)

    where coeffs are the actual mid-band DCT values and expected is the

    +1/-1 sequence for candidate_mark_id. An unmarked image gives score ~ 0.

    A correctly-marked image gives a positive peak clearly above noise.

    Threshold 0.015 is conservative; calibrate on your test set.

    Score for an incorrect mark_id is normally-distributed around 0 with

    stddev ~ 1/sqrt(n_coeffs), so for n_coeffs=1500, ~0.026. A correctly

    marked image typically scores > 0.03.

    """

    img = Image.open(io.BytesIO(image_bytes)).convert("RGB")

    ycbcr = img.convert("YCbCr")

    y = ycbcr.split()[0]

    y_arr = np.array(y, dtype=np.float64)

    D = _dct2(y_arr)

    coords = _pick_midband_indices(D.shape, n=n_coeffs)

    expected = _mark_to_sequence(candidate_mark_id, len(coords)).astype(np.float64)

    vals = np.array([D[i, j] for (i, j) in coords], dtype=np.float64)

    score = float(np.sum(vals * expected) / (np.sum(np.abs(vals)) + 1e-9))

    return (abs(score) >= threshold and score > 0), score

def perceptual_hash(image_bytes: bytes) -> str:

    """

    Perceptual hash (pHash) for fuzzy leak-match lookup.

    Uses imagehash. 64-bit output, hex-encoded.

    """

    import imagehash

    img = Image.open(io.BytesIO(image_bytes)).convert("RGB")

    return str(imagehash.phash(img))