"""
Code extraction and file download helpers.
"""

import io
import re
import zipfile
from typing import Optional
from pydantic import BaseModel

from fastapi import APIRouter
from fastapi.responses import StreamingResponse

from backend.services.code_extractor import extract_code_blocks

router = APIRouter()


class ExtractBody(BaseModel):
    markdown: str
    title: Optional[str] = None


@router.post("/extract")
def extract_files(body: ExtractBody):
    blocks = extract_code_blocks(body.markdown)
    return {"files": blocks}


@router.post("/download-zip")
def download_zip(body: ExtractBody):
    blocks = extract_code_blocks(body.markdown)
    if not blocks:
        return {"error": "No code blocks found"}

    # Keep LAST occurrence of each filename (latest version wins)
    file_map: dict[str, str] = {}
    for b in blocks:
        file_map[b["filename"]] = b["code"]

    if not file_map:
        return {"error": "No code blocks found"}

    # Build a safe zip filename from chat title
    raw_title = (body.title or "").strip()
    if not raw_title or raw_title == "New Chat":
        safe_title = "code"
    else:
        safe_title = re.sub(r'[^\w\s-]', '', raw_title).strip()
        safe_title = re.sub(r'[\s]+', '-', safe_title)[:60] or "code"
    zip_filename = f"{safe_title}.zip"

    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
        for name, code in file_map.items():
            zf.writestr(name, code)

    buf.seek(0)
    return StreamingResponse(
        buf,
        media_type="application/zip",
        headers={"Content-Disposition": f'attachment; filename="{zip_filename}"'},
    )