"""Video extraction and the small image primitives the pipeline shares."""
from __future__ import annotations

import subprocess
from typing import Iterator, Sequence

import numpy as np


def temporal_median3(frames: Sequence[np.ndarray]) -> np.ndarray:
    """Median of three same-shaped frames, as uint8."""
    if len(frames) != 3:
        raise ValueError("temporal_median3 wants exactly three frames")
    return np.median(np.stack(frames), axis=0).astype(np.uint8)


def box_blur3(gray: np.ndarray) -> np.ndarray:
    """3x3 box blur with edge clamping (no wraparound)."""
    if gray.ndim != 2:
        raise ValueError("box_blur3 wants a 2-D gray image")
    p = np.pad(gray.astype(np.float32), 1, mode="edge")
    total = (
        p[0:-2, 0:-2] + p[0:-2, 1:-1] + p[0:-2, 2:]
        + p[1:-1, 0:-2] + p[1:-1, 1:-1] + p[1:-1, 2:]
        + p[2:, 0:-2] + p[2:, 1:-1] + p[2:, 2:]
    )
    return np.rint(total / 9.0).astype(np.uint8)


def binarize(gray: np.ndarray, threshold: int = 127) -> np.ndarray:
    """True means black ink (the source is dark-on-light)."""
    return gray < threshold


def pack_lsb(mask: np.ndarray) -> bytes:
    """Pack rows LSB-first: bit 0 of a byte is the leftmost pixel."""
    if mask.ndim != 2 or mask.shape[1] % 8:
        raise ValueError("pack_lsb wants 2-D rows whose width is a multiple of 8")
    return np.packbits(mask.astype(bool), axis=1, bitorder="little").tobytes()


def extract_gray(video: str, start: float, duration: float, fps: int,
                 width: int, height: int) -> Iterator[np.ndarray]:
    """Yield (height, width) uint8 gray frames from ffmpeg."""
    vf = "fps=%s,scale=%d:%d:flags=lanczos,format=gray" % (
        fps, width, height)
    cmd = [
        "ffmpeg", "-v", "error",
        "-ss", str(start),
        "-i", video,
        "-t", str(duration),
        "-vf", vf,
        "-f", "rawvideo", "-",
    ]
    size = width * height
    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, bufsize=size * 4)
    try:
        while True:
            buf = proc.stdout.read(size)
            if len(buf) < size:
                break
            yield np.frombuffer(buf, dtype=np.uint8).reshape(height, width)
    finally:
        proc.stdout.close()
        proc.wait()
        if proc.returncode != 0:
            raise RuntimeError("ffmpeg failed with status %d" % proc.returncode)


def extract_binary(video: str, start: float, duration: float, fps: int,
                   width: int, height: int,
                   params: dict) -> Iterator[np.ndarray]:
    """Yield stylised binary frames: temporal median, blur, threshold.

    Exactly ``duration * fps`` frames are produced.  The first output is the
    first frame against itself; each later output is the median of the
    previous, current and next frame.
    """
    threshold = int(params.get("threshold", 127))
    want_blur = int(params.get("blur", 3))
    it = extract_gray(video, start, duration, fps, width, height)
    try:
        first = next(it)
    except StopIteration:
        return

    def render(median_frame: np.ndarray) -> np.ndarray:
        gray = box_blur3(median_frame) if want_blur >= 3 else median_frame
        return binarize(gray, threshold)

    a = first
    b = first
    yield render(temporal_median3([a, b, first]))
    for current in it:
        yield render(temporal_median3([a, b, current]))
        a, b = b, current