#!/usr/bin/env python3
import argparse
import os
import random
import subprocess
import tempfile

from PIL import Image, ImageChops, ImageStat


# ---------------------------------------------------------------------------
# Image generation
# ---------------------------------------------------------------------------

def generate_plasma(width: int, height: int, rng: random.Random, scale: int) -> Image.Image:
    """
    Multi-octave noise with 1/f-like spectrum (similar to ImageMagick plasma).
    Each octave is a tiny random image bilinearly upsampled to full size;
    amplitudes halve each octave so low frequencies dominate.

    `scale` is the coarsest feature size in pixels (absolute, image-size-
    independent).  Larger values → smoother image; 1 → pure pixel noise.
    """
    accum = [[0.0, 0.0, 0.0] for _ in range(width * height)]
    total = 0.0
    amplitude = 1.0
    # Clamp so we never generate patches larger than the image itself
    octave_scale = min(scale, max(width, height))

    while octave_scale >= 1:
        sw = max(2, width  // octave_scale + 2)
        sh = max(2, height // octave_scale + 2)
        patch = bytes(rng.randint(0, 255) for _ in range(sw * sh * 3))
        layer = Image.frombytes('RGB', (sw, sh), patch).resize(
            (width, height), Image.BILINEAR)
        for i, rgb in enumerate(layer.get_flattened_data()):  # type: ignore[attr-defined]
            for c in range(3):
                accum[i][c] += amplitude * rgb[c]
        total += amplitude
        amplitude *= 0.5
        octave_scale //= 2

    pixels = bytes(
        min(255, max(0, round(accum[i][c] / total)))
        for i in range(width * height)
        for c in range(3)
    )
    return Image.frombytes('RGB', (width, height), pixels)


# ---------------------------------------------------------------------------
# Simulation
# ---------------------------------------------------------------------------

def run_simulation(jpeg_path, ppm_path, vvp_path):
    """
    Decode jpeg_path into ppm_path with vvp sim.vvp and return the result as a PIL Image.
    """
    
    result = subprocess.run(
        ['vvp', vvp_path, f'+infile={jpeg_path}', f'+outfile={ppm_path}'],
        capture_output=True, text=True,
    )
    combined = result.stdout + result.stderr

    assert result.returncode == 0, f"simulation exited with code {result.returncode}:\n{combined}"
    assert 'ERROR:' not in combined, f"simulation reported error:\n{combined}"
    assert os.path.exists(ppm_path) and os.path.getsize(ppm_path) > 0, f"simulation produced no output:\n{combined}"

    return _read_ppm(ppm_path)


def _read_ppm(ppm_path):
    """Parse a P6 binary PPM file into a PIL Image."""
    with open(ppm_path, 'rb') as f:
        data = f.read()

    # Header written by testbench: "P6\n<w> <h>\n255\n"
    pos = 0
    lines = []
    while len(lines) < 3:
        nl = data.index(b'\n', pos)
        lines.append(data[pos:nl].decode().strip())
        pos = nl + 1

    assert lines[0] == 'P6', f"unexpected PPM magic {lines[0]!r}"

    w, h = map(int, lines[1].split())
    expected = w * h * 3
    actual = len(data) - pos
    assert actual == expected, f"PPM pixel data is {actual} bytes, expected {expected} ({w}x{h})"

    return Image.frombytes('RGB', (w, h), data[pos:])

# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description='Run jpeg_core testbench with randomly generated images.')
    parser.add_argument('--width',       type=int, default=32, help='Image width in pixels (default: 32)')
    parser.add_argument('--height',      type=int, default=32, help='Image height in pixels (default: 32)')
    parser.add_argument('--subsampling', action='store_true',
                        help='Use 4:2:0 chroma subsampling (default: 4:4:4)')
    parser.add_argument('--seed',    type=int, default=42,
                        help='Random seed (default: 42)')
    parser.add_argument('--quality', type=int, default=85,
                        help='JPEG quality 1-95 (default 85)')
    parser.add_argument('--scale', type=int, default=32,
                        help='Plasma coarsest feature size in pixels (default 32)')
    parser.add_argument('--save-input',  metavar='PATH', help='Save generated JPEG to this path')
    parser.add_argument('--save-output', metavar='PATH', help='Save RTL output PPM to this path')
    parser.add_argument('--vvp-path', metavar='PATH', default='jpeg_core_tb.vvp',
                        help='Path to compiled simulation (default: jpeg_core_tb.vvp)')
    args = parser.parse_args()

    mcu, subsampling_str = (16, '4:2:0') if args.subsampling else (8, '4:4:4')

    assert args.width  % mcu == 0, f"width {args.width} must be a multiple of {mcu} for {subsampling_str}"
    assert args.height % mcu == 0, f"height {args.height} must be a multiple of {mcu} for {subsampling_str}"
    assert os.path.exists(args.vvp_path), f"{args.vvp_path} not found - run 'make' first"

    print(f"input: width={args.width} height={args.height} YCbCr={subsampling_str} q={args.quality} seed={args.seed}")

    jpeg_path = args.save_input
    if jpeg_path is None:
        fd, jpeg_path = tempfile.mkstemp(suffix='.jpg', prefix='jpeg_core_tb_infile_')
        os.close(fd)

    ppm_path = args.save_output
    if ppm_path is None:
        fd, ppm_path = tempfile.mkstemp(suffix='.jpg', prefix='jpeg_core_tb_outfile_')
        os.close(fd)

    try:
        img_input = generate_plasma(args.width, args.height, random.Random(args.seed), args.scale)
        img_input.save(jpeg_path, format='JPEG', quality=args.quality, subsampling=2 if args.subsampling else 0)
        img_input = Image.open(jpeg_path).convert('RGB')  # re-open to decode saved jpeg image

        img_output = run_simulation(jpeg_path, ppm_path, args.vvp_path)

        assert img_output.size == img_input.size, f"dimension mismatch: img_input={img_input.width}x{img_input.height} img_output={img_output.width}x{img_output.height}"

        stat = ImageStat.Stat(ImageChops.difference(img_output, img_input))

        per_channel = {
            ch: {'max': stat.extrema[i][1], 'mean': stat.mean[i]}
            for i, ch in enumerate(['R', 'G', 'B'])
        }
        overall_max  = max(v['max']  for v in per_channel.values())
        overall_mean = sum(v['mean'] for v in per_channel.values()) / 3
        print(f"result: max_err={overall_max} mean_err={overall_mean:.2f}")

    finally:
        if args.save_input is None:
            os.unlink(jpeg_path)


if __name__ == '__main__':
    main()