# Patched ACE-Step REST API.
#
# Upstream's infer-api.py builds an 18-element positional `params`
# tuple and splats it into the pipeline. The pipeline's __call__
# actually expects 24 positional args — upstream's tuple is missing
# `format` at the head and the audio2audio/lora block at the tail.
# The result of the missing `format` head: every subsequent arg shifts
# left, the int `audio_duration` lands in `format`'s slot, and the
# pipeline blows up with "object of type 'int' has no len()" when it
# tries to validate the format string.
#
# This file copies upstream's behavior verbatim except for the params
# tuple, which now matches the canonical 24-arg pipeline signature
# inferred from the Gradio handler in acestep/ui/components.py.
# (Defaults for the new fields mirror Gradio's defaults: format=wav,
# audio2audio_enable=False, ref_audio_strength=0.5, no ref audio, no
# LoRA.)
from fastapi import FastAPI, HTTPException
from fastapi.responses import Response
from pydantic import BaseModel
from typing import List, Optional
import os
from acestep.pipeline_ace_step import ACEStepPipeline
import uuid

app = FastAPI(title="ACEStep Pipeline API (patched)")

# audio_format -> Content-Type. The wrapper writes the file in the
# format the upstream pipeline supports for the requested extension;
# we just need to label the bytes correctly on the wire.
_AUDIO_MIME = {
    "wav":  "audio/wav",
    "mp3":  "audio/mpeg",
    "flac": "audio/flac",
}


class ACEStepInput(BaseModel):
    checkpoint_path: str
    bf16: bool = True
    torch_compile: bool = False
    device_id: int = 0
    output_path: Optional[str] = None
    audio_duration: float
    prompt: str
    lyrics: str
    infer_step: int
    guidance_scale: float
    scheduler_type: str
    cfg_type: str
    omega_scale: float
    actual_seeds: List[int]
    guidance_interval: float
    guidance_interval_decay: float
    min_guidance_scale: float
    use_erg_tag: bool
    use_erg_lyric: bool
    use_erg_diffusion: bool
    oss_steps: List[int]
    guidance_scale_text: float = 0.0
    guidance_scale_lyric: float = 0.0
    # Optional new fields exposed for completeness — Gradio defaults.
    # lora_name_or_path: empty string trips HF Hub's repo-id validator
    # ("Repo id ... cannot ... be empty"), so default to None and only
    # forward it as a string when actually set.
    audio_format: str = "wav"
    audio2audio_enable: bool = False
    ref_audio_strength: float = 0.5
    ref_audio_input: Optional[str] = None
    lora_name_or_path: Optional[str] = None
    lora_weight: float = 1.0


class ACEStepOutput(BaseModel):
    status: str
    output_path: Optional[str]
    message: str


def initialize_pipeline(
    checkpoint_path: str, bf16: bool, torch_compile: bool, device_id: int
) -> ACEStepPipeline:
    os.environ["CUDA_VISIBLE_DEVICES"] = str(device_id)
    return ACEStepPipeline(
        checkpoint_dir=checkpoint_path,
        dtype="bfloat16" if bf16 else "float32",
        torch_compile=torch_compile,
    )


@app.post("/generate")
async def generate_audio(input_data: ACEStepInput):
    """Generate music; respond with the audio bytes inline.

    Pre-2026-05-11 this returned a JSON {output_path} reference to a
    file inside the container — useless to any external consumer
    since the path wasn't reachable from outside. Now streams the
    bytes back directly with the right Content-Type, and the file in
    /app/outputs/ is incidental bookkeeping (the pipeline writes
    there as part of its normal flow; we read it back for the wire).
    """
    try:
        model_demo = initialize_pipeline(
            input_data.checkpoint_path,
            input_data.bf16,
            input_data.torch_compile,
            input_data.device_id,
        )

        params = (
            input_data.audio_format,
            input_data.audio_duration,
            input_data.prompt,
            input_data.lyrics,
            input_data.infer_step,
            input_data.guidance_scale,
            input_data.scheduler_type,
            input_data.cfg_type,
            input_data.omega_scale,
            ", ".join(map(str, input_data.actual_seeds)),
            input_data.guidance_interval,
            input_data.guidance_interval_decay,
            input_data.min_guidance_scale,
            input_data.use_erg_tag,
            input_data.use_erg_lyric,
            input_data.use_erg_diffusion,
            ", ".join(map(str, input_data.oss_steps)),
            input_data.guidance_scale_text,
            input_data.guidance_scale_lyric,
            input_data.audio2audio_enable,
            input_data.ref_audio_strength,
            input_data.ref_audio_input,
            input_data.lora_name_or_path or "none",
            input_data.lora_weight,
        )

        ext = input_data.audio_format.lower()
        output_path = (
            input_data.output_path
            or f"/app/outputs/output_{uuid.uuid4().hex}.{ext}"
        )

        model_demo(*params, save_path=output_path)

        with open(output_path, "rb") as f:
            audio_bytes = f.read()

        return Response(
            content=audio_bytes,
            media_type=_AUDIO_MIME.get(ext, "application/octet-stream"),
            headers={
                # Surfaces the in-container path for debugging/log correlation
                # without putting it on the response body.
                "X-Output-Path": output_path,
            },
        )

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error generating audio: {str(e)}")


@app.get("/health")
async def health_check():
    return {"status": "healthy"}


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)