import io
import os
import json
from pathlib import Path
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse
from batch_processor import batch_embed, batch_extract
from watermark import embed_watermark, extract_watermark

app = FastAPI(title="AudioWatermarker API")

@app.post("/embed/")
async def api_embed(
    file: UploadFile = File(...),
    message: str = Form(...),
    password: str = Form(...)
):
    """Inserta watermark en un único archivo."""
    if file.filename is None:
        raise HTTPException(status_code=400, detail="No filename provided")
    suffix = Path(file.filename).suffix
    data = await file.read()
    # Guardar temp
    tmp_in = f"/tmp/in{suffix}"
    tmp_out = f"/tmp/out{suffix}"
    with open(tmp_in, "wb") as f: f.write(data)
    bits = embed_watermark(tmp_in, tmp_out, message.encode(), password)
    def iterfile():
        with open(tmp_out, "rb") as f:
            yield from f
    headers = {"X-Bits-Inserted": str(bits)}
    return StreamingResponse(iterfile(), media_type="audio/wav", headers=headers)

@app.post("/extract/")
async def api_extract(
    file: UploadFile = File(...),
    password: str = Form(...)
):
    """Extrae watermark de un archivo."""
    if file.filename is None:
        raise HTTPException(status_code=400, detail="No filename provided")
    suffix = Path(file.filename).suffix
    tmp_in = f"/tmp/in{suffix}"
    with open(tmp_in, "wb") as f: f.write(await file.read())
    try:
        payload = extract_watermark(tmp_in, password)
        return JSONResponse({"payload": payload.decode(errors="ignore")})
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

@app.post("/batch/embed/")
async def api_batch_embed(
    message: str = Form(...),
    password: str = Form(...),
    files: list[UploadFile] = File(...)
):
    """Batch embed: recibe múltiples archivos, devuelve ZIP con watermarked."""
    import zipfile, tempfile
    tmpdir = tempfile.mkdtemp()
    outdir = tempfile.mkdtemp()
    # guardar entrantes
    for f in files:
        if f.filename is None:
            continue
        path = os.path.join(tmpdir, f.filename)
        with open(path, "wb") as fh: fh.write(await f.read())
    results = batch_embed(tmpdir, outdir, message.encode(), password)
    # crear zip
    zip_path = os.path.join(tempfile.gettempdir(), "batch_embed.zip")
    with zipfile.ZipFile(zip_path, "w") as z:
        for fn in os.listdir(outdir):
            z.write(os.path.join(outdir, fn), arcname=fn)
    # retornar
    return StreamingResponse(
        open(zip_path, "rb"), media_type="application/zip",
        headers={"X-Results": json.dumps(results)}
    )

@app.post("/batch/extract/")
async def api_batch_extract(
    password: str = Form(...),
    files: list[UploadFile] = File(...)
):
    """Batch extract: retorna JSON con payloads o errores."""
    import tempfile, os
    tmpdir = tempfile.mkdtemp()
    for f in files:
        if f.filename is None:
            continue
        path = os.path.join(tmpdir, f.filename)
        with open(path, "wb") as fh: fh.write(await f.read())
    results = batch_extract(tmpdir, password)
    return JSONResponse(results)