#!/usr/bin/env python3
from __future__ import annotations

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

BLOCK_SIZE = 128
TYPE_RAW = 0
TYPE_BINARY = 1
TYPE_BAS_PROG = 2
TYPE_BAS_ARR = 3

def crc16(data: bytes) -> int:
    """Calculate CRC-16/CCITT-FALSE."""
    crc = 0xffff
    for value in data:
        crc ^= value << 8
        for _ in range(8):
            if crc & 0x8000:
                crc = (crc << 1) ^ 0x1021
            else:
                crc <<= 1
            crc &= 0xffff
    return crc


class PulseWriter:
    def __init__(
        self,
        sample_rate: int,
        t_short_us: float,
        t_long_us: float,
        amplitude: float,
    ) -> None:
        self.sample_rate = sample_rate
        self.short_samples = self._duration_samples(t_short_us)
        self.long_samples = self._duration_samples(t_long_us)
        self.high = round(32767 * amplitude)
        self.low = -self.high
        self.samples = bytearray()

        self.space(0.1)

    def space(self, time: float) -> None:
        self._level(self.low, int(time*self.sample_rate))

    def _duration_samples(self, duration_us: float) -> int:
        sample_count = round(duration_us * self.sample_rate / 1_000_000)
        if sample_count < 1:
            raise ValueError(
                f"pulse duration {duration_us:g} us is shorter than one sample"
            )
        return sample_count

    def _level(self, value: int, sample_count: int) -> None:
        self.samples.extend(struct.pack("<h", value) * sample_count)

    def short_pulse(self) -> None:
        self._level(self.high, self.short_samples)
        self._level(self.low, self.long_samples)

    def long_pulse(self) -> None:
        self._level(self.high, self.long_samples)
        self._level(self.low, self.short_samples)

    def marker(self) -> None:
        for _ in range(4):
            self._level(self.high, self.short_samples)
            self._level(self.low, self.short_samples)

    def bit(self, value: int) -> None:
        if value:
            self.long_pulse()
        else:
            self.short_pulse()

    def byte(self, value: int) -> None:
        # Transmit bytes most-significant bit first.
        for bit_number in range(7, -1, -1):
            self.bit((value >> bit_number) & 1)

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

    def leader(self) -> None:
        self.space(0.2)
        for _ in range(50):
            self._level(self.high, self.short_samples)
            self._level(self.low, self.short_samples)
        for _ in range(50):
            self._level(self.high, self.long_samples)
            self._level(self.low, self.long_samples)
        self.marker()
        for value in (0x55, 0x5A, 0xC3):
            self.byte(value)

    def block(self, type:int, data:bytes) -> None:
        self.leader()
        block_data = struct.pack("<B",type) + data[:BLOCK_SIZE].ljust(BLOCK_SIZE, b'\0')
        self.bytes(block_data)
        self.bytes(struct.pack("<H", crc16(block_data)))
        self.byte(0x00) # Make sure the last byte has its last edge

def header(pulses: PulseWriter, path:Path) -> None:
    fname = path.name
    fsize = os.path.getsize(path)
    data = (
        struct.pack("<HH", math.ceil(fsize/BLOCK_SIZE), fsize)
        + fname.encode()[:32].ljust(32, b' ')
    )
    pulses.block(TYPE_BINARY, data)
    pulses.block(TYPE_BINARY, data)

def raw(pulses: PulseWriter, data:bytes) -> None:
    pulses.block(TYPE_RAW, data)
    pulses.block(TYPE_RAW, data)

def main() -> None:
    parser = argparse.ArgumentParser(
        description="Generate the pulse-width encoded tape leader"
    )
    parser.add_argument("input", type=Path, help="input file to encode")
    parser.add_argument("output", type=Path, help="output WAV file")
    parser.add_argument("--sample-rate", type=int, default=48_000)
    parser.add_argument(
        "--t-short",
        type=float,
        default=250.0,
        metavar="MICROSECONDS",
        help="short pulse interval in microseconds (default: 250)",
    )
    parser.add_argument(
        "--t-long",
        type=float,
        default=500.0,
        metavar="MICROSECONDS",
        help="long pulse interval in microseconds (default: 500)",
    )
    parser.add_argument("--amplitude", type=float, default=0.75)
    args = parser.parse_args()

    if args.sample_rate <= 0:
        parser.error("--sample-rate must be positive")
    if args.t_short <= 0 or args.t_long <= 0:
        parser.error("pulse intervals must be positive")
    if args.t_short >= args.t_long:
        parser.error("--t-short must be less than --t-long")
    if not 0.0 <= args.amplitude <= 1.0:
        parser.error("--amplitude must be between 0 and 1")

    try:
        pulses = PulseWriter(
            args.sample_rate,
            args.t_short,
            args.t_long,
            args.amplitude,
        )
    except ValueError as error:
        parser.error(str(error))

    # create header block
    pulses.space(5.0)
    header(pulses, args.input)
    pulses.space(1.0)
    with open(args.input, "rb") as f:
        while True:
            data = f.read(BLOCK_SIZE)
            if not len(data):
                break
            raw(pulses, data)

    pulses.space(1.0)
    with wave.open(str(args.output), "wb") as wav:
        wav.setnchannels(1)
        wav.setsampwidth(2)
        wav.setframerate(args.sample_rate)
        wav.writeframes(pulses.samples)


if __name__ == "__main__":
    main()