backend/routers/status.py

from fastapi import APIRouter, HTTPException
from fastapi.responses import FileResponse
import os
import time
import json
import asyncio
from fastapi.responses import FileResponse, StreamingResponse
from core.state import job_queue, state
from core.database import db_manager

router = APIRouter(prefix="/api")

@router.get("/events/status/{task_id}")
async def events_status(task_id: str):
    async def event_generator():
        last_data = None
        while True:
            task = db_manager.get_task(task_id)
            if not task:
                yield f"data: {json.dumps({'error': 'Task not found'})}\n\n"
                break
                
            if task["status"] == "queued":
                # Calculate queue position
                pos = 1
                for job in list(job_queue.queue):
                    if isinstance(job, dict):
                        job_id = job.get("task_id")
                    else:
                        job_id = job[0]
                        
                    if job_id in state.cancelled_tasks:
                        continue
                        
                    if job_id == task_id:
                        break
                    pos += 1
                
                task["queue_position"] = pos
                task["eta_seconds"] = pos * 180 + (180 if state.active_task_id else 0)
                task["message"] = f"Waiting in queue... Position: {pos}"
            else:
                # Ensure processing/completed tasks don't have queue_position
                task.pop("queue_position", None)
                
            # Only send if changed to reduce bandwidth
            if str(task) != str(last_data):
                yield f"data: {json.dumps(task)}\n\n"
                last_data = task
                
            if task["status"] in ["completed", "failed", "cancelled"]:
                break
                
            await asyncio.sleep(1)

    return StreamingResponse(event_generator(), media_type="text/event-stream")

@router.post("/cancel/{task_id}")
async def cancel_task(task_id: str):
    task = db_manager.get_task(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")
        
    if task["status"] == "completed":
        raise HTTPException(status_code=400, detail="Task already completed")
        
    state.cancelled_tasks.add(task_id)
    
    # If the task is currently processing via a tracked subprocess, kill it
    if task_id in state.active_processes:
        process = state.active_processes[task_id]
        if process and process.poll() is None:
            try:
                process.terminate() # Or kill()
            except Exception:
                pass
                
    if task["status"] == "queued":
        db_manager.upsert_task(task_id, {
            "status": "cancelled",
            "message": "Task cancelled by user"
        })
        
    return {"status": "cancelled"}

@router.get("/status/{task_id}")
async def get_status(task_id: str):
    task = db_manager.get_task(task_id)
    if not task:
        # Check if it exists on disk (rehydration after server restart)
        final_stems_dir = os.path.join("temp_workdir", task_id, "final_stems")
        if os.path.exists(final_stems_dir):
            task = {
                "status": "completed",
                "progress": 100,
                "message": "Recovered from disk",
                "result_path": os.path.join("temp_workdir", task_id, f"custom_stems_{task_id}.zip"),
                "start_time": time.time(),
                "completed_time": time.time()
            }
            db_manager.upsert_task(task_id, task)
        else:
            raise HTTPException(status_code=404, detail="Task not found")
            
    if task["status"] == "queued":
        # Calculate queue position
        pos = 1
        for job in list(job_queue.queue):
            if isinstance(job, dict):
                job_id = job.get("task_id")
            else:
                job_id = job[0]
                
            if job_id in state.cancelled_tasks:
                continue
                
            if job_id == task_id:
                break
            pos += 1
        
        task["queue_position"] = pos
        task["eta_seconds"] = pos * 180 + (180 if state.active_task_id else 0)
        task["message"] = f"Waiting in queue... Position: {pos}"
    else:
        task.pop("queue_position", None)
        
    return task

@router.get("/download/{task_id}")
async def download_result(task_id: str):
    task = db_manager.get_task(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")
        
    if task["status"] != "completed" or not task.get("result_path"):
        raise HTTPException(status_code=400, detail="Task not completed yet")
        
    file_path = task["result_path"]
    filename = os.path.basename(file_path)
    
    media_type = "application/zip"
    if filename.endswith(".mp3"):
        media_type = "audio/mpeg"
    elif filename.endswith(".wav"):
        media_type = "audio/wav"
        
    return FileResponse(
        file_path,
        media_type=media_type,
        filename=filename
    )

from services.audio_service import package_custom_download
@router.get("/custom_download/{task_id}")
async def custom_download(
    task_id: str, 
    stems: str = "vocals", # comma separated
    format: str = "wav",
    chunked: str = "false",
    folder_name: str = "custom_stems"
):
    task = db_manager.get_task(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")
        
    if task["status"] != "completed":
        raise HTTPException(status_code=400, detail="Task not completed yet")
        
    requested_stems = [s.strip() for s in stems.split(",") if s.strip()]
    if not requested_stems:
        raise HTTPException(status_code=400, detail="No stems requested")
        
    try:
        zip_path = package_custom_download(
            task_id, 
            requested_stems, 
            output_format=format, 
            chunked=(chunked.lower() == "true"), 
            folder_name=folder_name
        )
        return FileResponse(
            zip_path,
            media_type="application/zip",
            filename=f"{folder_name}_{task_id}.zip"
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@router.get("/stream/{task_id}/{filename}")
async def stream_audio(task_id: str, filename: str):
    task = db_manager.get_task(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="Task not found")
        
    if task["status"] != "completed":
        raise HTTPException(status_code=400, detail="Task not completed yet")
        
    file_path = os.path.join("temp_workdir", task_id, "final_stems", filename)
    if not os.path.exists(file_path):
        raise HTTPException(status_code=404, detail="File not found")
        
    media_type = "application/octet-stream"
    if filename.endswith(".wav"):
        media_type = "audio/wav"
    elif filename.endswith(".mp3"):
        media_type = "audio/mpeg"
    elif filename.endswith(".json"):
        media_type = "application/json"
    elif filename.endswith(".srt") or filename.endswith(".txt") or filename.endswith(".csv"):
        media_type = "text/plain"
        
    return FileResponse(file_path, media_type=media_type)