#!/usr/bin/env python3
"""Generate short mechanical startup/shutdown ping MP3 files for LittleFS.

Outputs by default:
- data/ping_mech_up.mp3
- data/ping_mech_down.mp3

Requires:
- pip install lameenc
"""

from __future__ import annotations

import argparse
import math
import random
import struct
from pathlib import Path

import lameenc


SAMPLE_RATE = 44100


def damped_ping(t: float, f: float, click_amt: float, rng: random.Random) -> float:
    """A bright, metallic-ish pluck with fast decay and a tiny striker click."""
    env1 = math.exp(-t * 11.0)
    env2 = math.exp(-t * 17.0)
    env3 = math.exp(-t * 24.0)
    body = (
        0.72 * math.sin(2 * math.pi * f * t) * env1
        + 0.23 * math.sin(2 * math.pi * f * 2.05 * t) * env2
        + 0.12 * math.sin(2 * math.pi * f * 3.87 * t) * env3
    )
    click = rng.uniform(-1.0, 1.0) * click_amt * math.exp(-t * 120.0)
    return body + click


def synth_sequence(
    hit_times: list[float],
    freqs: list[float],
    length_s: float,
    fade_start_s: float,
    fade_len_s: float,
    gain: float,
    click_amt: float,
    seed: int,
) -> bytes:
    rng = random.Random(seed)
    n = int(SAMPLE_RATE * length_s)
    samples: list[int] = []

    for i in range(n):
        t = i / SAMPLE_RATE
        x = 0.0
        for ht, f in zip(hit_times, freqs):
            dt = t - ht
            if 0.0 <= dt <= 0.22:
                x += damped_ping(dt, f, click_amt, rng)

        fade = 1.0 if t < fade_start_s else max(0.0, (length_s - t) / fade_len_s)
        x *= gain * fade
        x = math.tanh(1.85 * x)

        s = int(max(-1.0, min(1.0, x)) * 32767)
        samples.append(s)

    return struct.pack("<" + "h" * len(samples), *samples)


def encode_mp3(pcm16: bytes, bitrate_kbps: int) -> bytes:
    enc = lameenc.Encoder()
    enc.set_in_sample_rate(SAMPLE_RATE)
    enc.set_channels(1)
    enc.set_bit_rate(bitrate_kbps)
    enc.set_quality(2)
    out = enc.encode(pcm16)
    out += enc.flush()
    return out


def write_mp3(path: Path, pcm16: bytes, bitrate_kbps: int) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    data = encode_mp3(pcm16, bitrate_kbps)
    path.write_bytes(data)
    print(f"Wrote {path} ({path.stat().st_size} bytes)")


def generate_up(out_dir: Path, bitrate_kbps: int, seed: int) -> None:
    hit_times = [0.00, 0.18, 0.36, 0.56, 0.78, 1.00]
    freqs = [740, 830, 930, 1040, 1160, 1290]
    pcm = synth_sequence(
        hit_times=hit_times,
        freqs=freqs,
        length_s=1.25,
        fade_start_s=1.10,
        fade_len_s=0.15,
        gain=0.28,
        click_amt=0.18,
        seed=seed,
    )
    write_mp3(out_dir / "ping_mech_up.mp3", pcm, bitrate_kbps)


def generate_down(out_dir: Path, bitrate_kbps: int, seed: int) -> None:
    hit_times = [0.00, 0.16, 0.33, 0.51, 0.71, 0.92]
    freqs = [1280, 1140, 1010, 900, 800, 710]
    pcm = synth_sequence(
        hit_times=hit_times,
        freqs=freqs,
        length_s=1.20,
        fade_start_s=1.04,
        fade_len_s=0.16,
        gain=0.28,
        click_amt=0.16,
        seed=seed + 1,
    )
    write_mp3(out_dir / "ping_mech_down.mp3", pcm, bitrate_kbps)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Generate mechanical ping startup/shutdown MP3 files")
    parser.add_argument(
        "--out-dir",
        default="data",
        help="Output directory for generated MP3 files (default: data)",
    )
    parser.add_argument(
        "--bitrate",
        type=int,
        default=128,
        help="MP3 bitrate in kbps (default: 128)",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed for click texture reproducibility (default: 42)",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    out_dir = Path(args.out_dir)
    generate_up(out_dir, args.bitrate, args.seed)
    generate_down(out_dir, args.bitrate, args.seed)


if __name__ == "__main__":
    main()