from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
import torch
import uvicorn
import io
import numpy as np

app = FastAPI(title="ACE-Step 1.5 API")

class GenerateRequest(BaseModel):
    prompt: str
    duration: int = 30
    steps: int = 10

@app.get("/health")
def health():
    return {"status": "ok", "model": "ACE-Step 1.5"}

@app.post("/generate")
def generate(req: GenerateRequest):
    try:
        print(f"Generating: {req.prompt} (duration={req.duration}s, steps={req.steps})")

        # Placeholder: integrate actual ACE-Step 1.5 model here
        def event_stream():
            yield f"data: {{\"status\": \"started\", \"prompt\": \"{req.prompt}\"}}\n\n"
            yield f"data: {{\"status\": \"generating\", \"progress\": 50}}\n\n"
            yield f"data: {{\"status\": \"done\", \"audio_url\": \"/audio/{hash(req.prompt)}\"}}\n\n"

        return StreamingResponse(event_stream(), media_type="text/event-stream")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/audio/{audio_id}")
def get_audio(audio_id: str):
    return {"audio_id": audio_id, "format": "wav"}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)