"""
Video watermarking for clawID.

Strategy:
  • Embed: watermark every Nth frame (default every 5 frames) with the SAME
    clawID UUID, using the invisible DWT-QIM algorithm.
  • Detect: sample up to `max_frames` frames, attempt decode on each; return
    the first successful result (majority-vote across samples for robustness).

Dependencies:
  imageio + imageio-ffmpeg  — install with:  pip install "imageio[ffmpeg]"

Supported formats (anything ffmpeg can read/write):
  mp4, mov, avi, mkv, webm, gif, ...

Usage::

    from clawid.video import embed_video, detect_video

    meta = embed_video(
        'input.mp4',
        'output_wm.mp4',
        metadata={'uid': 'creator123', 'platform': 'hawky.ai'},
    )
    print(meta['clawid'])

    result = detect_video('output_wm.mp4')
    if result:
        print(result['uid'])
"""

from __future__ import annotations

import os
import time
import uuid
from typing import Optional

from PIL import Image

from .core.payload import encode, to_bits
from .core.detect import _try_detect_qim, _try_detect_fm
from .algorithms.dwt_qim import embed_bits, DELTA, REPEAT

_PILLOW_EXTS = {'.gif', '.tiff', '.tif', '.apng'}

def _plugin_for(path: str) -> str:
    """Return 'pillow' for image-sequence formats, 'pyav' for video formats."""
    return 'pillow' if os.path.splitext(path)[1].lower() in _PILLOW_EXTS else 'pyav'


# How many frames to skip between watermarked frames (1 = every frame)
DEFAULT_FRAME_STRIDE = 5
# Max frames to sample during detection
DEFAULT_MAX_DETECT_FRAMES = 10
# Min size (pixels in smallest dimension) for watermarking a frame
_MIN_FRAME_DIM = 256


def embed_video(
    source: str,
    output_path: str,
    metadata: dict,
    frame_stride: int = DEFAULT_FRAME_STRIDE,
    delta: float = DELTA,
    repeat: int = REPEAT,
    **ffmpeg_kwargs,
) -> dict:
    """
    Embed a clawID watermark into a video by watermarking every Nth frame.

    All watermarked frames share the same 'clawid' UUID so detection can
    succeed from any single watermarked frame.

    Args:
        source:        Input video path (any ffmpeg-readable format)
        output_path:   Output video path
        metadata:      Creator metadata dict (uid, platform, …)
        frame_stride:  Watermark every frame_stride-th frame (default 5).
                       Lower = more robust but slower to encode.
        delta:         QIM step for invisible watermark
        repeat:        QIM repetition factor
        **ffmpeg_kwargs: Extra kwargs forwarded to imageio writer
                         (e.g. fps=24, quality=7, codec='h264')

    Returns:
        Finalised metadata dict (includes auto-added 'clawid' and 'ts').
    """
    try:
        import imageio.v3 as iio
    except ImportError:
        raise ImportError(
            "imageio[ffmpeg] is required for video watermarking. "
            "Install with: pip install 'imageio[ffmpeg]'"
        )

    # Finalise metadata
    metadata = dict(metadata)
    if 'clawid' not in metadata:
        metadata['clawid'] = str(uuid.uuid4())
    if 'ts' not in metadata:
        metadata['ts'] = int(time.time())

    payload_bytes = encode(metadata)
    bits = to_bits(payload_bytes)

    import numpy as np

    read_plugin = _plugin_for(source)
    write_plugin = _plugin_for(output_path)

    # Read input video properties
    props = iio.improps(source, plugin=read_plugin)
    fps = getattr(props, 'fps', 30) or 30

    def _process_frame(frame_idx, frame):
        pil_frame = Image.fromarray(frame)
        h, w = frame.shape[:2]
        if frame_idx % frame_stride == 0 and min(h, w) >= _MIN_FRAME_DIM:
            try:
                pil_frame = embed_bits(pil_frame, bits, delta=delta, repeat=repeat)
            except Exception:
                pass
        return np.array(pil_frame)

    if write_plugin == 'pillow':
        # Use PIL directly for lossless multi-frame formats (TIFF, etc.)
        # imageio's pillow writer doesn't support multi-frame TIFF reliably.
        out_pil = [
            Image.fromarray(_process_frame(i, f))
            for i, f in enumerate(iio.imiter(source, plugin=read_plugin))
        ]
        ext = os.path.splitext(output_path)[1].lower()
        if ext in ('.gif',):
            # GIF: palette-based — watermark won't survive, but preserve the write path
            out_pil[0].save(output_path, save_all=True, append_images=out_pil[1:], loop=0)
        else:
            # TIFF and others: lossless
            out_pil[0].save(output_path, save_all=True, append_images=out_pil[1:])
    else:
        writer_kwargs = {'fps': fps}
        writer_kwargs.update(ffmpeg_kwargs)
        with iio.imopen(output_path, 'w', plugin=write_plugin, **writer_kwargs) as writer:
            for frame_idx, frame in enumerate(iio.imiter(source, plugin=read_plugin)):
                writer.write(_process_frame(frame_idx, frame))

    return metadata


def detect_video(
    source: str,
    algorithm: str = 'auto',
    max_frames: int = DEFAULT_MAX_DETECT_FRAMES,
    frame_stride: int = DEFAULT_FRAME_STRIDE,
) -> Optional[dict]:
    """
    Detect and extract a clawID watermark from a video.

    Samples up to max_frames frames, tries to decode each one, and returns
    the first successfully decoded metadata dict. Uses majority vote if
    multiple frames decode to the same clawid (most common result wins).

    Args:
        source:        Video path
        algorithm:     'auto' | 'qim' | 'fm'  (same semantics as detect())
        max_frames:    Maximum number of frames to sample
        frame_stride:  Sample every frame_stride-th frame

    Returns:
        Metadata dict if a watermark is found in any sampled frame, else None.
    """
    try:
        import imageio.v3 as iio
    except ImportError:
        raise ImportError(
            "imageio[ffmpeg] is required for video detection. "
            "Install with: pip install 'imageio[ffmpeg]'"
        )

    from collections import Counter

    candidates: list[dict] = []

    read_plugin = _plugin_for(source)

    try:
        for frame_idx, frame in enumerate(iio.imiter(source, plugin=read_plugin)):
            if frame_idx % frame_stride != 0:
                continue

            pil_frame = Image.fromarray(frame)

            result = None
            if algorithm in ('auto', 'qim'):
                result = _try_detect_qim(pil_frame)
            if result is None and algorithm in ('auto', 'fm'):
                result = _try_detect_fm(pil_frame)

            if result is not None:
                candidates.append(result)

            if len(candidates) >= max_frames:
                break

    except Exception:
        pass

    if not candidates:
        return None

    # Majority vote on clawid UUID
    if len(candidates) == 1:
        return candidates[0]

    clawid_counts = Counter(c.get('clawid', '') for c in candidates)
    best_clawid = clawid_counts.most_common(1)[0][0]
    # Return the metadata from the first frame that had the winning clawid
    for c in candidates:
        if c.get('clawid') == best_clawid:
            return c

    return candidates[0]
