#!/usr/bin/env python3
"""Convert a Bad Apple source video into firmware-ready assets.

This tool requires ffmpeg in PATH. It produces:
- video.baa : 1-bit packed frames with custom header.
- audio.pdm : 1-bit sigma-delta encoded PCM data with custom header.
"""

import argparse
import math
import struct
import subprocess
import tempfile
from pathlib import Path

VIDEO_MAGIC = b"BAA\0"
AUDIO_MAGIC = b"BAP\0"


def run_ffmpeg(args):
    process = subprocess.run(args, check=False)
    if process.returncode != 0:
        raise RuntimeError(f"ffmpeg failed: {' '.join(args)}")


def pcm_to_pdm(samples):
    # First-order sigma-delta modulation.
    acc = 0
    threshold = 0
    bitstream = bytearray((len(samples) + 7) // 8)
    for idx, sample in enumerate(samples):
        acc += sample - threshold
        if acc >= 0:
            threshold = 32767
            bit = 1
        else:
            threshold = -32768
            bit = 0
        if bit:
            bitstream[idx // 8] |= 1 << (idx % 8)
    return bitstream


def encode_video(raw_path, width, height, frame_count, fps, threshold):
    stride_bits = width * height
    stride_bytes = (stride_bits + 7) // 8
    output = bytearray()
    output += VIDEO_MAGIC
    output += struct.pack('<H', width)
    output += struct.pack('<H', height)
    output += struct.pack('<H', fps)
    output += struct.pack('<I', frame_count)
    output += b"\x00\x00"

    frame_size = width * height
    with open(raw_path, 'rb') as fh:
        for frame_idx in range(frame_count):
            buf = fh.read(frame_size)
            if len(buf) != frame_size:
                raise RuntimeError(f"unexpected EOF at frame {frame_idx}")
            bits = bytearray(stride_bytes)
            for i, value in enumerate(buf):
                if value >= threshold:
                    bits[i // 8] |= 1 << (i % 8)
            output.extend(bits)
    return output


def encode_audio(raw_path, sample_rate, start_offset_samples=0):
    data = Path(raw_path).read_bytes()
    samples = struct.iter_unpack('<h', data)
    pcm = [value for (value,) in samples]
    if start_offset_samples > 0:
        pcm = pcm[start_offset_samples:]
    bitstream = pcm_to_pdm(pcm)

    header = bytearray()
    header += AUDIO_MAGIC
    header += struct.pack('<I', sample_rate)
    header += struct.pack('<B', 1)  # bits per sample
    header += struct.pack('<B', 1)  # mono
    header += b"\x00\x00"
    header += struct.pack('<I', len(pcm))
    return bytes(header) + bytes(bitstream)


parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--input', required=True, help='Path to Bad Apple video (any ffmpeg-supported format).')
parser.add_argument('--output-dir', required=True, help='Directory to place generated assets.')
parser.add_argument('--width', type=int, default=320)
parser.add_argument('--height', type=int, default=240)
parser.add_argument('--fps', type=int, default=30)
parser.add_argument('--sample-rate', type=int, default=44100)
parser.add_argument('--threshold', type=int, default=128, help='Luma threshold (0-255) for 1-bit conversion.')
parser.add_argument('--audio-start-delay', type=float, default=0.0, help='Seconds to delay audio relative to video.')
args = parser.parse_args()

output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)

with tempfile.TemporaryDirectory() as tmpdir:
    tmpdir = Path(tmpdir)
    raw_video = tmpdir / 'frames.raw'
    raw_audio = tmpdir / 'audio.raw'

    run_ffmpeg([
        'ffmpeg', '-y', '-i', args.input,
        '-vf', f'scale={args.width}:{args.height},fps={args.fps},format=gray',
        '-an', '-f', 'rawvideo', str(raw_video)
    ])

    run_ffmpeg([
        'ffmpeg', '-y', '-i', args.input,
        '-vn', '-ac', '1', '-ar', str(args.sample_rate), '-f', 's16le', str(raw_audio)
    ])

    frame_size = args.width * args.height
    video_size = raw_video.stat().st_size
    frame_count = video_size // frame_size
    if frame_count == 0:
        raise RuntimeError("ffmpeg produced zero frames - check input file")

    video_blob = encode_video(raw_video, args.width, args.height, frame_count, args.fps, args.threshold)
    audio_offset = int(args.audio_start_delay * args.sample_rate)
    audio_blob = encode_audio(raw_audio, args.sample_rate, audio_offset)

    (output_dir / 'video.baa').write_bytes(video_blob)
    (output_dir / 'audio.pdm').write_bytes(audio_blob)

    print(f"Generated {frame_count} frames @ {args.width}x{args.height} and audio track {args.sample_rate} Hz")