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

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

class FSKWriter:
    def __init__(self, sample_rate: int, c_low: int, c_high: int, f_low: float, f_high: float, amplitude: float):
        self.sample_rate = sample_rate
        self.c_low = c_low
        self.c_high = c_high
        self.f_low = f_low
        self.f_high = f_high
        self.amplitude = amplitude
        self.samples_per_low = (sample_rate * c_low) // f_low
        self.samples_per_high = (sample_rate * c_high) // f_high
        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(int(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) -> None:
        self._tone(self.f_high, (200*self.sample_rate)//self.f_high)
        self._tone(self.f_low, (200*self.sample_rate)//self.f_low)
        for _ in range(8):
            self.bit(1)
            self.bit(0)
        self.byte(0x0f)

    def bit(self, value: int) -> None:
        if value:
            self._tone(self.f_high, self.samples_per_high)
        else:
            self._tone(self.f_low, self.samples_per_low)

    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 block(self, nr: int, data: bytes) -> None:
        self.leader()
        self.byte(nr)
        l = len(data)
        self.byte(l>>8)
        self.byte(l&0xff)
        self.bytes(data)
        self.byte(self.chsum(data))

    def chsum(self, data: bytes) -> int:
        cs = 0
        for d in data:
            cs += int(d)
        return cs


def main() -> None:
    parser = argparse.ArgumentParser(description="Encode a binary file as two tone FSK WAV")
    parser.add_argument("input", type=Path, help="Binary input file")
    parser.add_argument("output", type=Path, help="output WAVl default <input>.wav")
    parser.add_argument("--sample-rate", type=int, default=48000)
    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("--low-cycles", type=int, default=2, help="low tone cycle count for a 0")
    parser.add_argument("--high-cycles", type=int, default=4, help="high tone cycle count for a 1")
    parser.add_argument("--amplitude", type=float, default=0.75)
    parser.add_argument("--block-size", type=int, default=64)
    args = parser.parse_args()

    output_path = args.output or args.input.with_name(args.input.stem + ".wav")

    tones = FSKWriter(args.sample_rate, args.low_cycles, args.high_cycles, args.low, args.high, args.amplitude)
    tones.block(0, args.input.name.encode())


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

if __name__ == "__main__":
    main()