"""Download endpoint for split results."""

from fastapi import APIRouter, HTTPException
from fastapi.responses import FileResponse
from zipfile import ZipFile
import os

from backend.config import settings
from backend.services.task_manager import task_manager
from backend.models.response import TaskStatus

router = APIRouter(prefix="/api", tags=["download"])


def _prepare_download(task_id: str):
    """Common logic to prepare and return the ZIP file."""
    if not task_manager.has_task(task_id):
        raise HTTPException(status_code=404, detail=f"Task {task_id} not found")

    status = task_manager.get_status(task_id)
    if status["status"] != TaskStatus.DONE:
        raise HTTPException(
            status_code=400,
            detail=f"Task {task_id} is not complete. Current status: {status['status']}"
        )

    output_dir = settings.temp_dir / task_id / "output"
    if not output_dir.exists() or not any(output_dir.iterdir()):
        raise HTTPException(status_code=404, detail="No output files found")

    zip_path = settings.temp_dir / task_id / "splits.zip"
    # Create ZIP if not already exists (or overwrite)
    with ZipFile(zip_path, "w") as zipf:
        for file_path in output_dir.iterdir():
            if file_path.is_file():
                zipf.write(file_path, arcname=file_path.name)

    return zip_path


@router.get("/download/{task_id}")
async def download_results(task_id: str):
    zip_path = _prepare_download(task_id)
    return FileResponse(
        path=zip_path,
        media_type="application/zip",
        filename=f"{task_id}_splits.zip",
        headers={"Content-Disposition": f"attachment; filename={task_id}_splits.zip"}
    )


@router.get("/download/{task_id}/splits.zip")
async def download_results_as_zip(task_id: str):
    """Alias endpoint that provides a .zip suffix for easier curl usage."""
    zip_path = _prepare_download(task_id)
    return FileResponse(
        path=zip_path,
        media_type="application/zip",
        filename="splits.zip",
        headers={"Content-Disposition": "attachment; filename=splits.zip"}
    )