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", )