import io
from pathlib import Path

from pydub import AudioSegment

from backend.key_sound_map import (
    MANIFEST,
    marker_release_time,
    resolve_marker_samples,
    sample_path,
)
from backend.models import ExportRequest

SAMPLE_RATE = 48000


def _load_sample(path: Path) -> AudioSegment:
    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:
    if request.switch not in MANIFEST:
        raise FileNotFoundError(f"Unknown switch: {request.switch}")

    duration_ms = int(request.duration * 1000)
    base = AudioSegment.silent(duration=duration_ms, frame_rate=SAMPLE_RATE)

    for marker in sorted(request.markers, key=lambda m: m.time):
        press_key, release_key = resolve_marker_samples(
            request.switch,
            marker.key,
            marker.code,
        )
        press_sample = _load_sample(sample_path(request.switch, "press", press_key))
        release_sample = _load_sample(sample_path(request.switch, "release", release_key))

        press_ms = int(marker.time * 1000)
        if press_ms < duration_ms:
            base = base.overlay(press_sample, position=press_ms)

        release_time = marker_release_time(marker.time, marker.release_time)
        release_ms = int(release_time * 1000)
        if release_ms < duration_ms:
            base = base.overlay(release_sample, position=release_ms)

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