import click
import os

from acestep.pipeline_ace_step import ACEStepPipeline
from acestep.data_sampler import DataSampler


def sample_data(json_data):
    return (
        json_data["audio_duration"],
        json_data["prompt"],
        json_data["lyrics"],
        json_data["infer_step"],
        json_data["guidance_scale"],
        json_data["scheduler_type"],
        json_data["cfg_type"],
        json_data["omega_scale"],
        ", ".join(map(str, json_data["actual_seeds"])),
        json_data["guidance_interval"],
        json_data["guidance_interval_decay"],
        json_data["min_guidance_scale"],
        json_data["use_erg_tag"],
        json_data["use_erg_lyric"],
        json_data["use_erg_diffusion"],
        ", ".join(map(str, json_data["oss_steps"])),
        json_data["guidance_scale_text"] if "guidance_scale_text" in json_data else 0.0,
        (
            json_data["guidance_scale_lyric"]
            if "guidance_scale_lyric" in json_data
            else 0.0
        ),
    )


@click.command()
@click.option(
    "--checkpoint_path", type=str, default="", help="Path to the checkpoint directory"
)
@click.option("--bf16", type=bool, default=True, help="Whether to use bfloat16")
@click.option(
    "--torch_compile", type=bool, default=False, help="Whether to use torch compile"
)
@click.option(
    "--cpu_offload", type=bool, default=False, help="Whether to use CPU offloading (only load current stage's model to GPU)"
)
@click.option(
    "--overlapped_decode", type=bool, default=False, help="Whether to use overlapped decoding (run dcae and vocoder using sliding windows)"
)
@click.option("--device_id", type=int, default=0, help="Device ID to use")
@click.option("--output_path", type=str, default=None, help="Path to save the output")
# --- generation parameters -------------------------------------------------
# Without --prompt the script keeps its original behaviour and generates from a
# random example in examples/input_params.
@click.option(
    "--prompt",
    type=str,
    default=None,
    help="Style tags or description, comma separated. If omitted, a random example is used.",
)
@click.option("--lyrics", type=str, default=None, help="Lyrics, with [verse]/[chorus] tags.")
@click.option(
    "--lyrics_file",
    type=click.Path(exists=True, dir_okay=False),
    default=None,
    help="Read lyrics from a UTF-8 file. Takes precedence over --lyrics.",
)
@click.option("--duration", type=float, default=60.0, help="Audio duration in seconds (-1 for random).")
@click.option("--steps", type=int, default=60, help="Number of inference steps.")
@click.option("--guidance_scale", type=float, default=15.0, help="Guidance scale.")
@click.option(
    "--scheduler",
    type=click.Choice(["euler", "heun", "pingpong"]),
    default="euler",
    help="Scheduler type.",
)
@click.option(
    "--cfg_type",
    type=click.Choice(["cfg", "apg", "cfg_star"]),
    default="apg",
    help="CFG type. apg is recommended.",
)
@click.option("--omega_scale", type=float, default=10.0, help="Granularity scale.")
@click.option("--seed", type=str, default=None, help="Manual seeds, comma separated. Random if omitted.")
@click.option(
    "--format",
    "audio_format",
    type=click.Choice(["wav", "mp3", "ogg", "flac"]),
    default="wav",
    help="Output audio format.",
)
def main(
    checkpoint_path,
    bf16,
    torch_compile,
    cpu_offload,
    overlapped_decode,
    device_id,
    output_path,
    prompt,
    lyrics,
    lyrics_file,
    duration,
    steps,
    guidance_scale,
    scheduler,
    cfg_type,
    omega_scale,
    seed,
    audio_format,
):
    os.environ["CUDA_VISIBLE_DEVICES"] = str(device_id)

    model_demo = ACEStepPipeline(
        checkpoint_dir=checkpoint_path,
        dtype="bfloat16" if bf16 else "float32",
        torch_compile=torch_compile,
        cpu_offload=cpu_offload,
        overlapped_decode=overlapped_decode
    )

    if lyrics_file is not None:
        with open(lyrics_file, "r", encoding="utf-8") as f:
            lyrics = f.read()

    if prompt is None:
        # No prompt given: keep the original behaviour and use a random example.
        data_sampler = DataSampler()
        json_data = data_sampler.sample()
        (
            audio_duration,
            prompt,
            sampled_lyrics,
            infer_step,
            sampled_guidance_scale,
            scheduler_type,
            sampled_cfg_type,
            sampled_omega_scale,
            manual_seeds,
            guidance_interval,
            guidance_interval_decay,
            min_guidance_scale,
            use_erg_tag,
            use_erg_lyric,
            use_erg_diffusion,
            oss_steps,
            guidance_scale_text,
            guidance_scale_lyric,
        ) = sample_data(json_data)
        # Explicit options still win over the sampled example.
        lyrics = lyrics if lyrics is not None else sampled_lyrics
        manual_seeds = seed if seed is not None else manual_seeds
    else:
        audio_duration = duration
        lyrics = lyrics if lyrics is not None else "[instrumental]"
        infer_step = steps
        sampled_guidance_scale = guidance_scale
        scheduler_type = scheduler
        sampled_cfg_type = cfg_type
        sampled_omega_scale = omega_scale
        manual_seeds = seed
        guidance_interval = 0.5
        guidance_interval_decay = 0.0
        min_guidance_scale = 3.0
        use_erg_tag = True
        use_erg_lyric = True
        use_erg_diffusion = True
        oss_steps = None
        guidance_scale_text = 0.0
        guidance_scale_lyric = 0.0

    click.echo(f"prompt: {prompt}")
    click.echo(f"duration: {audio_duration}s, steps: {infer_step}, seeds: {manual_seeds}")

    model_demo(
        format=audio_format,
        audio_duration=audio_duration,
        prompt=prompt,
        lyrics=lyrics,
        infer_step=infer_step,
        guidance_scale=sampled_guidance_scale,
        scheduler_type=scheduler_type,
        cfg_type=sampled_cfg_type,
        omega_scale=sampled_omega_scale,
        manual_seeds=manual_seeds,
        guidance_interval=guidance_interval,
        guidance_interval_decay=guidance_interval_decay,
        min_guidance_scale=min_guidance_scale,
        use_erg_tag=use_erg_tag,
        use_erg_lyric=use_erg_lyric,
        use_erg_diffusion=use_erg_diffusion,
        oss_steps=oss_steps,
        guidance_scale_text=guidance_scale_text,
        guidance_scale_lyric=guidance_scale_lyric,
        save_path=output_path,
    )


if __name__ == "__main__":
    main()