#!/usr/bin/env python3
"""
z80_tape_encode.py

Encode a binary file as a mono PCM WAV using robust two-tone Manchester FSK.

Defaults:
- 48 kHz sample rate
- 600 bit/s
- 2 kHz / 4 kHz tones
- Manchester FSK:
    0 -> low tone for half a bit, then high tone
    1 -> high tone for half a bit, then low tone
- 1 second 4 kHz leader
- 64 bytes of 0x55 preamble
- 256-byte CRC-protected blocks
- CRC-16/CCITT-FALSE
"""

from __future__ import annotations

import argparse
import math
import struct
import wave
from pathlib import Path

SYNC = b"\xD3\x91"
PREAMBLE_BYTE = 0x55


def crc16_ccitt_false(data: bytes) -> int:
    crc = 0xFFFF
    for byte in data:
        crc ^= byte << 8
        for _ in range(8):
            crc = ((crc << 1) ^ 0x1021) & 0xFFFF if (crc & 0x8000) else (crc << 1) & 0xFFFF
    return crc


class FSKWriter:
    def __init__(self, sample_rate: int, baud: int, f_low: float, f_high: float, amplitude: float):
        if sample_rate % (2 * baud) != 0:
            raise ValueError("sample_rate must be divisible by 2*baud")
        self.sample_rate = sample_rate
        self.baud = baud
        self.f_low = f_low
        self.f_high = f_high
        self.amplitude = amplitude
        self.samples_per_half_bit = sample_rate // (2 * baud)
        self.phase = 0.0
        self.samples = bytearray()

    def _tone(self, frequency: float, sample_count: int) -> None:
        phase_step = 2.0 * math.pi * frequency / self.sample_rate
        scale = int(32767 * self.amplitude)
        phase = self.phase
        out = self.samples

        for _ in range(sample_count):
            sample = int(round(scale * math.sin(phase)))
            out.extend(struct.pack("<h", sample))
            phase += phase_step
            if phase >= 2.0 * math.pi:
                phase -= 2.0 * math.pi

        self.phase = phase

    def leader(self, seconds: float) -> None:
        self._tone(self.f_high, round(seconds * self.sample_rate))

    def bit(self, value: int) -> None:
        # Manchester-FSK:
        # 0 = low -> high
        # 1 = high -> low
        if value:
            first, second = self.f_high, self.f_low
        else:
            first, second = self.f_low, self.f_high

        self._tone(first, self.samples_per_half_bit)
        self._tone(second, self.samples_per_half_bit)

    def byte(self, value: int) -> None:
        # MSB first
        for bit_no in range(7, -1, -1):
            self.bit((value >> bit_no) & 1)

    def bytes(self, data: bytes) -> None:
        for value in data:
            self.byte(value)


def make_block(sequence: int, payload: bytes) -> bytes:
    header = struct.pack("<HH", sequence & 0xFFFF, len(payload))
    crc = crc16_ccitt_false(header + payload)
    return SYNC + header + payload + struct.pack("<H", crc)


def make_eof_block(sequence: int) -> bytes:
    header = struct.pack("<HH", sequence & 0xFFFF, 0)
    crc = crc16_ccitt_false(header)
    return SYNC + header + struct.pack("<H", crc)


def encode(
    input_path: Path,
    output_path: Path,
    *,
    sample_rate: int = 48_000,
    baud: int = 600,
    f_low: float = 2_000.0,
    f_high: float = 4_000.0,
    leader_seconds: float = 1.0,
    preamble_bytes: int = 64,
    block_size: int = 256,
    amplitude: float = 0.75,
) -> None:
    data = input_path.read_bytes()

    if not 0 < amplitude <= 1.0:
        raise ValueError("amplitude must be > 0 and <= 1")
    if not 1 <= block_size <= 0xFFFF:
        raise ValueError("block_size must be in 1..65535")

    writer = FSKWriter(sample_rate, baud, f_low, f_high, amplitude)

    writer.leader(leader_seconds)
    writer.bytes(bytes([PREAMBLE_BYTE]) * preamble_bytes)

    sequence = 0
    for start in range(0, len(data), block_size):
        payload = data[start:start + block_size]
        writer.bytes(make_block(sequence, payload))
        sequence = (sequence + 1) & 0xFFFF

    writer.bytes(make_eof_block(sequence))

    with wave.open(str(output_path), "wb") as wav:
        wav.setnchannels(1)
        wav.setsampwidth(2)
        wav.setframerate(sample_rate)
        wav.writeframes(writer.samples)

    duration = len(writer.samples) / 2 / sample_rate

    print(f"Input:        {input_path}")
    print(f"Output:       {output_path}")
    print(f"Input bytes:  {len(data)}")
    print(f"Bit rate:     {baud} bit/s")
    print(f"Tones:        {f_low:g} Hz / {f_high:g} Hz")
    print(f"Sample rate:  {sample_rate} Hz")
    print(f"Half-bit:     {writer.samples_per_half_bit} samples")
    print(f"Block size:   {block_size} bytes")
    print(f"WAV duration: {duration:.2f} s")


def main() -> None:
    parser = argparse.ArgumentParser(description="Encode a binary file as Manchester two-tone FSK WAV.")
    parser.add_argument("input", type=Path, help="binary input file")
    parser.add_argument("output", type=Path, nargs="?", help="output WAV; default: INPUT.wav")
    parser.add_argument("--baud", type=int, default=600)
    parser.add_argument("--low", type=float, default=2000.0, help="low tone in Hz")
    parser.add_argument("--high", type=float, default=4000.0, help="high tone in Hz")
    parser.add_argument("--sample-rate", type=int, default=48000)
    parser.add_argument("--leader", type=float, default=1.0, help="leader seconds")
    parser.add_argument("--preamble", type=int, default=64, help="number of 0x55 preamble bytes")
    parser.add_argument("--block-size", type=int, default=64)
    parser.add_argument("--amplitude", type=float, default=0.75)

    args = parser.parse_args()
    output = args.output or args.input.with_name(args.input.name + ".wav")

    encode(
        args.input,
        output,
        sample_rate=args.sample_rate,
        baud=args.baud,
        f_low=args.low,
        f_high=args.high,
        leader_seconds=args.leader,
        preamble_bytes=args.preamble,
        block_size=args.block_size,
        amplitude=args.amplitude,
    )


if __name__ == "__main__":
    main()