#!/usr/bin/env python3
"""Play or convert the firmware's length-prefixed raw Opus prompt files."""

from __future__ import annotations

import argparse
from pathlib import Path
import shutil
import struct
import subprocess
import sys
import tempfile


MAX_PACKET_BYTES = 2_048
OGG_SIGNATURE = b"OggS"
OPUS_CLOCK_HZ = 48_000
PROMPT_FRAME_MS = 60
OGG_STREAM_SERIAL = 0x54514341


class PromptError(RuntimeError):
    """Raised when a prompt file does not match the firmware format."""


def read_prompt_packets(path: Path) -> list[bytes]:
    packets: list[bytes] = []
    with path.open("rb") as stream:
        while True:
            header = stream.read(4)
            if not header:
                break
            if len(header) != 4:
                raise PromptError("truncated 4-byte packet header")
            packet_size = struct.unpack(">I", header)[0]
            if packet_size == 0 or packet_size > MAX_PACKET_BYTES:
                raise PromptError(f"invalid Opus packet size: {packet_size}")
            packet = stream.read(packet_size)
            if len(packet) != packet_size:
                raise PromptError("truncated Opus packet payload")
            packets.append(packet)
    if not packets:
        raise PromptError("prompt contains no Opus packets")
    return packets


def ogg_crc(data: bytes) -> int:
    crc = 0
    for value in data:
        crc ^= value << 24
        for _ in range(8):
            crc = ((crc << 1) ^ 0x04C11DB7) & 0xFFFFFFFF if crc & 0x80000000 else (crc << 1) & 0xFFFFFFFF
    return crc


def packet_lacing(packet_size: int) -> bytes:
    segments = bytearray()
    remaining = packet_size
    while remaining >= 255:
        segments.append(255)
        remaining -= 255
    segments.append(remaining)
    return bytes(segments)


def make_ogg_page(packet: bytes, flags: int, granule: int, sequence: int) -> bytes:
    lacing = packet_lacing(len(packet))
    if len(lacing) > 255:
        raise PromptError("Opus packet requires too many Ogg segments")
    header = bytearray()
    header.extend(OGG_SIGNATURE)
    header.extend((0, flags))
    header.extend(struct.pack("<Q", granule))
    header.extend(struct.pack("<I", OGG_STREAM_SERIAL))
    header.extend(struct.pack("<I", sequence))
    header.extend(b"\x00\x00\x00\x00")
    header.append(len(lacing))
    header.extend(lacing)
    page = header + packet
    struct.pack_into("<I", page, 22, ogg_crc(page))
    return bytes(page)


def write_ogg(path: Path, packets: list[bytes]) -> None:
    opus_head = b"OpusHead" + struct.pack("<BBHIhB", 1, 1, 0, 16_000, 0, 0)
    vendor = b"TQ_CatN prompt player"
    opus_tags = b"OpusTags" + struct.pack("<I", len(vendor)) + vendor + struct.pack("<I", 0)
    samples_per_packet = OPUS_CLOCK_HZ * PROMPT_FRAME_MS // 1_000

    with path.open("wb") as output:
        output.write(make_ogg_page(opus_head, flags=0x02, granule=0, sequence=0))
        output.write(make_ogg_page(opus_tags, flags=0, granule=0, sequence=1))
        granule = 0
        for index, packet in enumerate(packets):
            granule += samples_per_packet
            flags = 0x04 if index == len(packets) - 1 else 0
            output.write(make_ogg_page(packet, flags=flags, granule=granule, sequence=index + 2))


def convert_prompt(input_path: Path, output_path: Path) -> None:
    if not input_path.is_file():
        raise PromptError(f"file not found: {input_path}")
    if input_path.read_bytes()[:4] == OGG_SIGNATURE:
        shutil.copyfile(input_path, output_path)
        return
    write_ogg(output_path, read_prompt_packets(input_path))


def play_prompt(input_path: Path, ffplay: str) -> None:
    executable = shutil.which(ffplay)
    if executable is None:
        raise PromptError("ffplay is not installed; run: brew install ffmpeg")
    with tempfile.TemporaryDirectory(prefix="tq_prompt_player_") as temp_dir:
        ogg_path = Path(temp_dir) / f"{input_path.stem}.ogg"
        convert_prompt(input_path, ogg_path)
        subprocess.run(
            [executable, "-hide_banner", "-loglevel", "warning", "-nodisp", "-autoexit", str(ogg_path)],
            check=True,
        )


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("input", type=Path, help="firmware prompt .opus file")
    parser.add_argument("--output", type=Path, help="write standard Ogg Opus instead of playing")
    parser.add_argument("--ffplay", default="ffplay", help="ffplay executable name or path")
    args = parser.parse_args()

    try:
        if args.output is not None:
            convert_prompt(args.input, args.output)
            print(f"Created standard Ogg Opus: {args.output}")
        else:
            play_prompt(args.input, args.ffplay)
    except (PromptError, OSError, subprocess.CalledProcessError) as error:
        print(f"Error: {error}", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())