#!/usr/bin/env python3
"""Convert an audio file into the device's length-prefixed raw Opus format.

The runtime prompt player does not consume an Ogg Opus file directly. It
expects a sequence of records with this layout:

    uint32_be packet_size
    uint8_t   raw_opus_packet[packet_size]

FFmpeg is used only for decoding/resampling and Opus encoding. This script
then removes the Ogg container and writes the framing expected by the device.
"""

from __future__ import annotations

import argparse
import os
from pathlib import Path
import shutil
import struct
import subprocess
import sys
import tempfile
from typing import BinaryIO, Iterable, Iterator


SAMPLE_RATE_HZ = 16_000
CHANNELS = 1
FRAME_DURATION_MS = 60
DEFAULT_BITRATE_KBPS = 24
MAX_PACKET_BYTES = 2_048
OGG_CAPTURE_PATTERN = b"OggS"


class ConversionError(RuntimeError):
    """Raised when conversion input or output violates the prompt contract."""


def _read_exact(stream: BinaryIO, size: int, description: str) -> bytes:
    data = stream.read(size)
    if len(data) != size:
        raise ConversionError(f"truncated {description}")
    return data


def iter_ogg_packets(path: Path) -> Iterator[bytes]:
    """Yield packets from the single logical Ogg stream generated by FFmpeg."""

    packet = bytearray()
    stream_serial: int | None = None
    expected_sequence: int | None = None

    with path.open("rb") as stream:
        while True:
            capture = stream.read(4)
            if capture == b"":
                break
            if capture != OGG_CAPTURE_PATTERN:
                raise ConversionError("invalid Ogg capture pattern")

            fixed_header = _read_exact(stream, 23, "Ogg page header")
            version = fixed_header[0]
            header_type = fixed_header[1]
            serial = struct.unpack_from("<I", fixed_header, 10)[0]
            sequence = struct.unpack_from("<I", fixed_header, 14)[0]
            segment_count = fixed_header[22]
            if version != 0:
                raise ConversionError(f"unsupported Ogg version: {version}")
            if stream_serial is None:
                stream_serial = serial
                expected_sequence = sequence
            if serial != stream_serial:
                raise ConversionError("multiple Ogg logical streams are unsupported")
            if expected_sequence is not None and sequence != expected_sequence:
                raise ConversionError(
                    f"unexpected Ogg page sequence: {sequence}, "
                    f"expected {expected_sequence}"
                )
            expected_sequence = sequence + 1

            continued = (header_type & 0x01) != 0
            if continued != bool(packet):
                raise ConversionError("invalid Ogg continued-packet flag")

            lacing = _read_exact(stream, segment_count, "Ogg segment table")
            page_body = _read_exact(stream, sum(lacing), "Ogg page body")
            offset = 0
            for segment_size in lacing:
                packet.extend(page_body[offset : offset + segment_size])
                offset += segment_size
                if segment_size < 255:
                    yield bytes(packet)
                    packet.clear()

    if packet:
        raise ConversionError("truncated Ogg packet at end of file")


def opus_packet_duration_us(packet: bytes) -> int:
    """Return an Opus packet duration using the TOC rules from RFC 6716."""

    if not packet:
        raise ConversionError("empty Opus packet")

    toc = packet[0]
    config = toc >> 3
    if config >= 16:
        frame_duration_us = 2_500 << (config & 0x03)
    elif config >= 12:
        frame_duration_us = 10_000 << (config & 0x01)
    elif (config & 0x03) == 0x03:
        frame_duration_us = 60_000
    else:
        frame_duration_us = 10_000 << (config & 0x03)

    frame_code = toc & 0x03
    if frame_code == 0:
        frame_count = 1
    elif frame_code in (1, 2):
        frame_count = 2
    else:
        if len(packet) < 2:
            raise ConversionError("truncated Opus packet frame-count byte")
        frame_count = packet[1] & 0x3F
        if frame_count == 0:
            raise ConversionError("invalid zero-frame Opus packet")

    duration_us = frame_duration_us * frame_count
    if duration_us > 120_000:
        raise ConversionError(
            f"Opus packet duration exceeds 120 ms: {duration_us / 1000:g} ms"
        )
    return duration_us


def validate_raw_prompt_packets(packets: Iterable[bytes]) -> tuple[int, int]:
    packet_count = 0
    payload_bytes = 0
    expected_duration_us = FRAME_DURATION_MS * 1_000
    for index, packet in enumerate(packets, start=1):
        if len(packet) > MAX_PACKET_BYTES:
            raise ConversionError(
                f"packet {index} is {len(packet)} bytes; maximum is "
                f"{MAX_PACKET_BYTES}"
            )
        duration_us = opus_packet_duration_us(packet)
        if duration_us != expected_duration_us:
            raise ConversionError(
                f"packet {index} duration is {duration_us / 1000:g} ms; "
                f"expected {FRAME_DURATION_MS} ms"
            )
        packet_count += 1
        payload_bytes += len(packet)

    if packet_count == 0:
        raise ConversionError("no audio packets were produced")
    return packet_count, payload_bytes


def read_length_prefixed_packets(path: Path) -> Iterator[bytes]:
    with path.open("rb") as stream:
        while True:
            header = stream.read(4)
            if header == b"":
                break
            if len(header) != 4:
                raise ConversionError("truncated prompt packet header")
            packet_size = struct.unpack(">I", header)[0]
            if packet_size == 0 or packet_size > MAX_PACKET_BYTES:
                raise ConversionError(
                    f"invalid prompt packet size: {packet_size}"
                )
            yield _read_exact(stream, packet_size, "prompt packet payload")


def write_length_prefixed_packets(path: Path, packets: Iterable[bytes]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary_path: Path | None = None
    try:
        with tempfile.NamedTemporaryFile(
            mode="wb", prefix=f".{path.name}.", dir=path.parent, delete=False
        ) as output:
            temporary_path = Path(output.name)
            for packet in packets:
                output.write(struct.pack(">I", len(packet)))
                output.write(packet)
        os.replace(temporary_path, path)
        temporary_path = None
    finally:
        if temporary_path is not None:
            temporary_path.unlink(missing_ok=True)


def encode_ogg_opus(
    ffmpeg: str, input_path: Path, output_path: Path, bitrate_kbps: int
) -> None:
    command = [
        ffmpeg,
        "-hide_banner",
        "-loglevel",
        "error",
        "-y",
        "-i",
        str(input_path),
        "-map",
        "0:a:0",
        "-vn",
        "-map_metadata",
        "-1",
        "-ac",
        str(CHANNELS),
        "-ar",
        str(SAMPLE_RATE_HZ),
        "-c:a",
        "libopus",
        "-application",
        "voip",
        "-frame_duration",
        str(FRAME_DURATION_MS),
        "-b:a",
        f"{bitrate_kbps}k",
        "-vbr",
        "off",
        "-compression_level",
        "10",
        str(output_path),
    ]
    try:
        subprocess.run(command, check=True)
    except FileNotFoundError as error:
        raise ConversionError(f"ffmpeg executable not found: {ffmpeg}") from error
    except subprocess.CalledProcessError as error:
        raise ConversionError(f"ffmpeg failed with exit code {error.returncode}") from error


def convert(
    input_path: Path,
    output_path: Path,
    ffmpeg: str,
    bitrate_kbps: int,
    overwrite: bool,
) -> tuple[int, int]:
    if not input_path.is_file():
        raise ConversionError(f"input file does not exist: {input_path}")
    if output_path.exists() and not overwrite:
        raise ConversionError(
            f"output already exists: {output_path}; pass --force to replace it"
        )
    if not 6 <= bitrate_kbps <= 510:
        raise ConversionError("bitrate must be between 6 and 510 kbps")

    resolved_ffmpeg = shutil.which(ffmpeg)
    if resolved_ffmpeg is None:
        raise ConversionError(
            f"ffmpeg executable not found: {ffmpeg}; install FFmpeg or pass "
            "--ffmpeg /absolute/path/to/ffmpeg"
        )

    with tempfile.TemporaryDirectory(prefix="tq_prompt_opus_") as temp_dir:
        ogg_path = Path(temp_dir) / "encoded.ogg"
        encode_ogg_opus(
            resolved_ffmpeg, input_path, ogg_path, bitrate_kbps
        )
        ogg_packets = list(iter_ogg_packets(ogg_path))

    if len(ogg_packets) < 3 or not ogg_packets[0].startswith(b"OpusHead") or not ogg_packets[1].startswith(b"OpusTags"):
        raise ConversionError("FFmpeg output is missing OpusHead/OpusTags packets")

    audio_packets = ogg_packets[2:]
    packet_count, payload_bytes = validate_raw_prompt_packets(audio_packets)
    write_length_prefixed_packets(output_path, audio_packets)
    return packet_count, payload_bytes


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Convert MP3 or another FFmpeg-supported audio file into the "
            "length-prefixed raw Opus prompt format used by TQ_CatN."
        )
    )
    parser.add_argument("input", nargs="?", type=Path, help="source audio file")
    parser.add_argument("output", nargs="?", type=Path, help="destination .opus file")
    parser.add_argument(
        "--bitrate-kbps",
        type=int,
        default=DEFAULT_BITRATE_KBPS,
        help=f"Opus bitrate in kbps (default: {DEFAULT_BITRATE_KBPS})",
    )
    parser.add_argument(
        "--ffmpeg",
        default="ffmpeg",
        help="FFmpeg executable name or absolute path (default: ffmpeg)",
    )
    parser.add_argument(
        "--force", action="store_true", help="replace an existing output file"
    )
    parser.add_argument(
        "--validate-only",
        type=Path,
        metavar="PROMPT_OPUS",
        help="validate an existing device prompt without invoking FFmpeg",
    )
    args = parser.parse_args()
    if args.validate_only is None and (args.input is None or args.output is None):
        parser.error("input and output are required unless --validate-only is used")
    if args.validate_only is not None and (args.input is not None or args.output is not None):
        parser.error("input/output cannot be combined with --validate-only")
    return args


def main() -> int:
    args = parse_args()
    try:
        if args.validate_only is not None:
            packet_count, payload_bytes = validate_raw_prompt_packets(
                read_length_prefixed_packets(args.validate_only)
            )
            output_path = args.validate_only
            action = "Validated"
        else:
            packet_count, payload_bytes = convert(
                args.input,
                args.output,
                args.ffmpeg,
                args.bitrate_kbps,
                args.force,
            )
            output_path = args.output
            action = "Created"
    except (ConversionError, OSError) as error:
        print(f"error: {error}", file=sys.stderr)
        return 1

    print(
        f"{action} {output_path}: {packet_count} packets, "
        f"{payload_bytes} Opus bytes, {SAMPLE_RATE_HZ} Hz mono, "
        f"{FRAME_DURATION_MS} ms frames"
    )
    return 0


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