#!/usr/bin/env python3
"""Generate a song with the native Apple-Silicon MiniMax-Music3 MLX port."""

from __future__ import annotations

import argparse
import wave
from pathlib import Path

import numpy as np

from minimax_mlx_model import SAMPLE_RATE, MiniMaxMusic3MlxPipeline


def bounded_float(minimum: float, maximum: float):
    def parse(value: str) -> float:
        parsed = float(value)
        if not minimum <= parsed <= maximum:
            raise argparse.ArgumentTypeError(f"must be between {minimum:g} and {maximum:g}")
        return parsed

    return parse


def bounded_int(minimum: int, maximum: int):
    def parse(value: str) -> int:
        parsed = int(value)
        if not minimum <= parsed <= maximum:
            raise argparse.ArgumentTypeError(f"must be between {minimum} and {maximum}")
        return parsed

    return parse


def write_wav(path: Path, audio: np.ndarray) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    pcm = np.clip(audio.T, -1.0, 1.0)
    pcm = np.round(pcm * 32_767).astype("<i2")
    with wave.open(str(path), "wb") as output:
        output.setnchannels(2)
        output.setsampwidth(2)
        output.setframerate(SAMPLE_RATE)
        output.writeframes(pcm.tobytes())


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--prompt", required=True, help="Musical direction/caption")
    lyrics = parser.add_mutually_exclusive_group(required=True)
    lyrics.add_argument("--lyrics", help="Lyrics text; use [Instrumental] for no vocals")
    lyrics.add_argument("--lyrics-file", type=Path, help="UTF-8 text file containing lyrics")
    parser.add_argument("--seconds", type=bounded_float(10, 300), default=60.0)
    parser.add_argument("--steps", type=bounded_int(1, 30), default=30)
    parser.add_argument("--seed", type=bounded_int(0, 2**31 - 1), default=7)
    parser.add_argument("--model-dir", type=Path, default=Path(__file__).resolve().parent)
    parser.add_argument("--output", type=Path, default=Path("song.wav"))
    args = parser.parse_args()

    lyrics_text = (
        args.lyrics_file.read_text(encoding="utf-8") if args.lyrics_file else args.lyrics
    )
    assert lyrics_text is not None

    def progress(stage: int, message: str) -> None:
        print(f"[{stage}/5] {message}", flush=True)

    pipeline = MiniMaxMusic3MlxPipeline(args.model_dir, progress)
    audio = pipeline.generate(
        args.prompt.strip(),
        lyrics_text.strip(),
        args.seconds,
        args.steps,
        args.seed,
        progress,
    )
    progress(5, f"Writing {args.output}…")
    write_wav(args.output, audio)
    print(args.output.resolve())


if __name__ == "__main__":
    main()