import shutil
import subprocess
from pathlib import Path
from fastapi import FastAPI, UploadFile, File, BackgroundTasks, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

UPLOAD_DIR = Path("uploads")
OUTPUT_DIR = Path("separated")

UPLOAD_DIR.mkdir(exist_ok=True)
OUTPUT_DIR.mkdir(exist_ok=True)

app.mount("/stems", StaticFiles(directory=OUTPUT_DIR), name="stems")


def separate_audio_task(file_path: Path, track_name: str):
    try:
        command = [
            "demucs",
            "-n", "htdemucs",
            "-o", str(OUTPUT_DIR),
            str(file_path)
        ]
        subprocess.run(command, check=True)
    except subprocess.CalledProcessError as e:
        print(f"Error during separation: {e}")
    finally:
        if file_path.exists():
            file_path.unlink()



@app.post("/upload")
async def upload_audio(background_tasks: BackgroundTasks, file: UploadFile = File(...)):
    extension = file.filename.split(".")[-1].lower()
    if extension not in ["mp3", "wav", "m4a", "flac"]:
        raise HTTPException(status_code=400, detail="Unsupported audio format.")

    track_name = Path(file.filename).stem
    temp_file_path = UPLOAD_DIR / file.filename
    
    with temp_file_path.open("wb") as buffer:
        shutil.copyfileobj(file.file, buffer)

    background_tasks.add_task(separate_audio_task, temp_file_path, track_name)
    base_url = f"/stems/htdemucs/{track_name}"
    
    return {
        "status": "processing",
        "track_name": track_name,
        "stems": {
            "vocals": f"{base_url}/vocals.wav",
            "drums": f"{base_url}/drums.wav",
            "bass": f"{base_url}/bass.wav",
            "other": f"{base_url}/other.wav"
        }
    }



@app.get("/status/{track_name}")
async def check_status(track_name: str):
    expected_folder = OUTPUT_DIR / "htdemucs" / track_name
    if expected_folder.exists() and (expected_folder / "vocals.wav").exists():
        return {"status": "completed"}
    return {"status": "processing"}