app/main.py

import csv
import json
import os
from dataclasses import asdict
from pathlib import Path
from tempfile import NamedTemporaryFile

from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles

from app.jobs import jobs, start_job


BASE_DIR = Path(__file__).resolve().parent.parent
STATIC_DIR = BASE_DIR / "static"
RUNS_DIR = BASE_DIR / "runs"
DEFAULT_AI_MODE_URL = "https://www.google.com/search?udm=50"
DEFAULT_BATCH_SIZE = int(os.getenv("PARALLEL_TABS", "10"))

RUNS_DIR.mkdir(exist_ok=True)

app = FastAPI(title="Bank Statement AI Mode Extractor")
app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        origin.strip()
        for origin in os.getenv("CORS_ORIGINS", "*").split(",")
        if origin.strip()
    ],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")


@app.get("/")
def index():
    if os.getenv("API_ONLY", "false").strip().lower() in {"1", "true", "yes"}:
        return FileResponse(STATIC_DIR / "backend.html")
    return FileResponse(STATIC_DIR / "index.html")


@app.get("/api/info")
def api_info() -> dict[str, str]:
    return {
        "name": "Bank Statement AI Mode Extractor API",
        "health": "/api/health",
        "create_job": "POST /api/jobs with multipart fields pdf and optional pwd",
        "job_status": "/api/jobs/{job_id}",
        "latest_browser_view": "/api/jobs/{job_id}/browser-view/latest",
        "download_json": "/api/jobs/{job_id}/download",
        "download_csv": "/api/jobs/{job_id}/download.csv",
    }


@app.get("/api/health")
def health() -> dict[str, str]:
    return {"status": "ok"}


@app.post("/api/jobs")
async def create_job(
    pdf: UploadFile = File(...),
    pwd: str = Form(""),
) -> dict[str, str]:
    if not pdf.filename or not pdf.filename.lower().endswith(".pdf"):
        raise HTTPException(status_code=400, detail="Please upload a PDF file.")
    with NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:
        content = await pdf.read()
        temp_file.write(content)
        temp_path = Path(temp_file.name)

    job = start_job(
        temp_path,
        RUNS_DIR,
        ai_mode_url=DEFAULT_AI_MODE_URL,
        batch_size=max(1, min(DEFAULT_BATCH_SIZE, 10)),
        pdf_password=pwd.strip() or None,
    )
    return {"job_id": job.id}


@app.get("/api/jobs/{job_id}")
def get_job(job_id: str) -> dict:
    job = jobs.get(job_id)
    if not job:
        raise HTTPException(status_code=404, detail="Job not found.")
    return asdict(job)


@app.get("/api/jobs/{job_id}/browser-view/latest")
def latest_browser_view(job_id: str) -> FileResponse:
    job = jobs.get(job_id)
    if not job or not job.latest_screenshot:
        raise HTTPException(status_code=404, detail="Browser screenshot is not ready.")
    return FileResponse(job.latest_screenshot, media_type="image/png")


@app.get("/api/jobs/{job_id}/browser-view/page/{page_number}")
def page_browser_view(job_id: str, page_number: int) -> FileResponse:
    job = jobs.get(job_id)
    if not job or not job.run_dir:
        raise HTTPException(status_code=404, detail="Job not found.")
    screenshot = Path(job.run_dir) / "browser-view" / f"page-{page_number:04d}.png"
    if not screenshot.exists():
        raise HTTPException(status_code=404, detail="Browser screenshot is not ready.")
    return FileResponse(screenshot, media_type="image/png")


@app.get("/api/jobs/{job_id}/download")
def download_job(job_id: str) -> FileResponse:
    job = jobs.get(job_id)
    if not job or not job.output_file:
        raise HTTPException(status_code=404, detail="Output file is not ready.")
    return FileResponse(
        job.output_file,
        media_type="application/json",
        filename=f"bank-statement-{job_id}.json",
    )


@app.get("/api/jobs/{job_id}/download.csv")
def download_job_csv(job_id: str) -> FileResponse:
    job = jobs.get(job_id)
    if not job or not job.output_file:
        raise HTTPException(status_code=404, detail="Output file is not ready.")

    json_path = Path(job.output_file)
    csv_path = json_path.with_suffix(".csv")
    payload = json.loads(json_path.read_text(encoding="utf-8"))
    columns = payload.get("columns", ["date", "description", "voucher_type", "amount", "closing"])
    rows = payload.get("data", [])

    with csv_path.open("w", newline="", encoding="utf-8-sig") as csv_file:
        writer = csv.DictWriter(csv_file, fieldnames=columns, extrasaction="ignore")
        writer.writeheader()
        for row in rows:
            writer.writerow(row)

    return FileResponse(
        csv_path,
        media_type="text/csv",
        filename=f"bank-statement-{job_id}.csv",
    )