import io
from pathlib import Path

from pydub import AudioSegment

from backend.models import ExportRequest, SwitchType

SAMPLE_RATE = 48000
SAMPLES_DIR = Path(__file__).resolve().parent.parent / "assets" / "samples"

SWITCH_FILES: dict[SwitchType, str] = {
    "cherry-mx-blue": "cherry-mx-blue.wav",
    "cherry-mx-red": "cherry-mx-red.wav",
    "cherry-mx-brown": "cherry-mx-brown.wav",
}


def _load_sample(switch: SwitchType) -> AudioSegment:
    path = SAMPLES_DIR / SWITCH_FILES[switch]
    if not path.exists():
        raise FileNotFoundError(f"Sample not found: {path}")
    sample = AudioSegment.from_file(path)
    return sample.set_frame_rate(SAMPLE_RATE).set_channels(2)


def export_keyboard_audio(request: ExportRequest) -> bytes:
    duration_ms = int(request.duration * 1000)
    base = AudioSegment.silent(duration=duration_ms, frame_rate=SAMPLE_RATE)
    click = _load_sample(request.switch)

    for marker in sorted(request.markers, key=lambda m: m.time):
        position_ms = int(marker.time * 1000)
        if position_ms < duration_ms:
            base = base.overlay(click, position=position_ms)

    buffer = io.BytesIO()
    base.export(buffer, format="wav")
    return buffer.getvalue()