main.py

"""
FindLEI  –  FastAPI backend
=============================
Endpoints
---------
POST /api/upload              Upload Excel; returns job_id + preview
POST /api/process/{job_id}    Start async LEI batch check
GET  /api/stream/{job_id}     SSE stream (real-time progress)
GET  /api/status/{job_id}     Poll-based status + results
GET  /api/download/{job_id}   Download enriched Excel
GET  /                        Serve frontend (static/index.html)
"""

import asyncio
import json
import logging
import time
import uuid
import zipfile
import io
import olefile
import os
import hmac
import hashlib
import concurrent.futures
import unicodedata
import threading
import secrets

from collections import defaultdict
from contextlib import asynccontextmanager
from functools import partial
from pathlib import Path
from typing import Dict, Optional

from fastapi import BackgroundTasks, FastAPI, HTTPException, Request, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse, Response, StreamingResponse
from fastapi.staticfiles import StaticFiles

from lei_checker import check_lei_batch
from log_config import setup_logging
from metrics import (
    active_jobs, http_requests_total, job_duration_seconds,
    jobs_total, lei_duration_seconds, leis_checked_total, metrics_response,
)

from excel_handler import (
    ExcelReadError, UnsupportedFileError, OversizedFileError,
    CorruptFileError, ParserError,
    read_lei_from_excel, write_results_to_excel
)

# ── Constants ──────────────────────────────────────────────────────────────────
RATE_LIMIT_UPLOADS      = int(os.getenv("RATE_LIMIT_UPLOADS", "5"))
RATE_LIMIT_WINDOW       = int(os.getenv("RATE_LIMIT_WINDOW", "60"))
MAX_SHEETS              = int(os.getenv("MAX_SHEETS", "5"))
MAX_CELLS               = int(os.getenv("MAX_CELLS", "50000"))
MAX_LEIS                = int(os.getenv("MAX_LEIS", "5000"))
MAX_JOBS_PER_IP         = int(os.getenv("MAX_JOBS_PER_IP", "2"))
MAX_JOBS_GLOBAL         = int(os.getenv("MAX_JOBS_GLOBAL", "25"))
MAX_ZIP_ENTRIES         = int(os.getenv("MAX_ZIP_ENTRIES", "200"))
MAX_DECOMPRESSED_BYTES  = int(os.getenv("MAX_DECOMPRESSED_BYTES", str(25 * 1024 * 1024)))
MAX_COMPRESSION_RATIO   = int(os.getenv("MAX_COMPRESSION_RATIO", "20"))
MAX_UPLOAD_BYTES        = int(os.getenv("MAX_UPLOAD_BYTES", str(10 * 1024 * 1024)))
JOB_SECRET              = os.getenv("JOB_SECRET", "change-me-in-production")
PARSE_TIMEOUT_SECONDS   = int(os.getenv("PARSE_TIMEOUT_SECONDS", "30"))
PROCESS_TIMEOUT_SECONDS = int(os.getenv("PROCESS_TIMEOUT_SECONDS", "300"))
JOB_TTL_SECONDS         = int(os.getenv("JOB_TTL_SECONDS", "3600"))
CLEANUP_INTERVAL        = int(os.getenv("CLEANUP_INTERVAL_SECONDS", "300"))
METRICS_TOKEN           = os.getenv("METRICS_TOKEN", "")
METRICS_ALLOWED_IPS     = set(filter(None, os.getenv("METRICS_ALLOWED_IPS", "127.0.0.1,::1").split(",")))
_PRODUCTION             = os.getenv("ENVIRONMENT", "development").lower() == "production"

_parse_executor = concurrent.futures.ThreadPoolExecutor(max_workers=4, thread_name_prefix="findlei-parse")
_upload_log: dict = defaultdict(list)
_upload_lock = threading.Lock()

# ── Logging ───────────────────────────────────────────────────────────────────
setup_logging()
logger = logging.getLogger("findlei.main")

# ── In-memory job store ───────────────────────────────────────────────────────
jobs: Dict[str, dict] = {}

# ── Lifespan ──────────────────────────────────────────────────────────────────
@asynccontextmanager
async def lifespan(app):
    asyncio.create_task(_periodic_cleanup())
    yield

# ── App ───────────────────────────────────────────────────────────────────────
app = FastAPI(
    title="FindLEI API",
    description="LEI batch lookup for banking compliance",
    version="1.0.0",
    lifespan=lifespan,
    docs_url=None if _PRODUCTION else "/docs",
    redoc_url=None if _PRODUCTION else "/redoc",
    openapi_url=None if _PRODUCTION else "/openapi.json",
)

@app.middleware("http")
async def limit_upload_size(request: Request, call_next):
    if request.method == "POST" and "/api/upload" in request.url.path:
        content_length = request.headers.get("content-length")
        if content_length and int(content_length) > MAX_UPLOAD_BYTES:
            from fastapi.responses import JSONResponse
            return JSONResponse(status_code=413, content={"detail": "File too large (max 10 MB)"})
    return await call_next(request)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.middleware("http")
async def _instrument(request: Request, call_next):
    start = time.perf_counter()
    response = await call_next(request)
    http_requests_total.labels(
        method=request.method,
        path=request.url.path,
        status_code=str(response.status_code),
    ).inc()
    logger.debug(
        "http request",
        extra={
            "method":      request.method,
            "path":        request.url.path,
            "status":      response.status_code,
            "duration_ms": round((time.perf_counter() - start) * 1000, 1),
        },
    )
    return response

# ── Security Headers ──────────────────────────────────────────────────────────

@app.middleware("http")
async def add_security_headers(request: Request, call_next):
    response = await call_next(request)
    response.headers["X-Content-Type-Options"]  = "nosniff"
    response.headers["X-Frame-Options"]         = "DENY"
    response.headers["Referrer-Policy"]         = "strict-origin-when-cross-origin"
    response.headers["Permissions-Policy"]      = "geolocation=(), camera=(), microphone=()"
    response.headers["Content-Security-Policy"] = (
        "default-src 'self'; "
        "script-src 'self' 'unsafe-inline'; "
        "style-src 'self' 'unsafe-inline' https://fonts.googleapis.com; "
        "font-src 'self' https://fonts.gstatic.com; "
        "img-src 'self' data:; "
        "connect-src 'self'; "
        "frame-ancestors 'none';"
    )
    if _PRODUCTION:
        response.headers["Strict-Transport-Security"] = "max-age=63072000; includeSubDomains; preload"
    return response

# ── Helpers ───────────────────────────────────────────────────────────────────

def _sanitize_filename(filename: str) -> str:
    if not filename:
        return "upload.xlsx"
    filename = unicodedata.normalize("NFKC", filename)
    filename = filename.replace("/", "_").replace("\\", "_").replace("\x00", "")
    filename = "".join(c for c in filename if unicodedata.category(c) not in ("Cc", "Cf") and c not in "\r\n")
    if len(filename) > 200:
        ext = filename.rsplit(".", 1)[-1] if "." in filename else "xlsx"
        filename = f"{filename[:195]}.{ext}"
    return filename.strip() or "upload.xlsx"

def _check_upload_rate_limit(ip: str) -> bool:
    now = time.time()
    with _upload_lock:
        timestamps = _upload_log.get(ip, [])
        _upload_log[ip] = [t for t in timestamps if now - t < RATE_LIMIT_WINDOW]
        if len(_upload_log[ip]) >= RATE_LIMIT_UPLOADS:
            return False
        _upload_log[ip].append(now)
        return True

def _count_active_jobs_for_ip(ip: str) -> int:
    return sum(
        1 for j in jobs.values()
        if j.get("client_ip") == ip and j.get("status") in ("pending", "processing")
    )

def _count_active_jobs_global() -> int:
    return sum(
        1 for j in jobs.values()
        if j.get("status") in ("pending", "processing")
    )

def _get_job(job_id: str) -> dict:
    if job_id not in jobs:
        raise HTTPException(status_code=404, detail="Job not found")
    return jobs[job_id]

def _validate_file_magic(content: bytes, suffix: str) -> bool:
    if suffix in {".xlsx", ".xlsm", ".ods"}:
        return content[:4] == b"PK\x03\x04"
    if suffix == ".xls":
        return content[:8] == b"\xD0\xCF\x11\xE0\xA1\xB1\x1A\xE1"
    return False

def validate_modern_excel_archive(content: bytes, suffix: str) -> bool:
    if suffix not in {".xlsx", ".xlsm", ".ods"}:
        return True
    try:
        with zipfile.ZipFile(io.BytesIO(content)) as zf:
            names = zf.namelist()
            if suffix in {".xlsx", ".xlsm"}:
                return "[Content_Types].xml" in names and "xl/workbook.xml" in names
            if suffix == ".ods":
                return "mimetype" in names and "content.xml" in names
    except zipfile.BadZipFile:
        return False
    return False

def _check_decompression_bomb(content: bytes, suffix: str) -> Optional[str]:
    if suffix not in {".xlsx", ".xlsm", ".ods"}:
        return None
    try:
        with zipfile.ZipFile(io.BytesIO(content)) as zf:
            entries = zf.infolist()
            if len(entries) > MAX_ZIP_ENTRIES:
                return "too_many_entries"
            total_uncompressed = 0
            for entry in entries:
                total_uncompressed += entry.file_size
                if total_uncompressed > MAX_DECOMPRESSED_BYTES:
                    return "too_large"
            for entry in entries:
                if entry.compress_size > 0:
                    ratio = entry.file_size / entry.compress_size
                    if ratio > MAX_COMPRESSION_RATIO:
                        return "high_ratio"
    except zipfile.BadZipFile:
        return "bad_zip"
    return None

def _validate_xls_ole2(content: bytes) -> bool:
    try:
        ole = olefile.OleFileIO(io.BytesIO(content))
        has_workbook = ole.exists("Workbook") or ole.exists("Book")
        ole.close()
        return has_workbook
    except Exception:
        return False

def _make_job_token(job_id: str) -> str:
    return hmac.new(
        JOB_SECRET.encode(),
        job_id.encode(),
        hashlib.sha256
    ).hexdigest()

def _verify_job_token(job_id: str, token: str) -> bool:
    expected = _make_job_token(job_id)
    return hmac.compare_digest(expected, token)

def _verify_ownership(job_id: str, job_token: Optional[str]) -> None:
    if not job_token or not _verify_job_token(job_id, job_token):
        raise HTTPException(status_code=403, detail="Invalid or missing job token.")

def _cleanup_expired_jobs():
    now = time.time()
    expired = [
        job_id for job_id, job in list(jobs.items())
        if now - job.get("created_at", now) > JOB_TTL_SECONDS
    ]
    for job_id in expired:
        job = jobs.get(job_id)
        if job:
            job["original_bytes"] = b""
            job["leis"]           = []
            job["results"]        = []
            job["error_msg"]      = ""
            job["status"]         = "expired"
        jobs.pop(job_id, None)
    if expired:
        logger.info("job_completed", extra={"job_id": job_id, "count": len(job["results"])})

async def _periodic_cleanup():
    while True:
        await asyncio.sleep(CLEANUP_INTERVAL)
        _cleanup_expired_jobs()


def _size_bucket(size_bytes: int) -> str:
    """Return human-readable size bucket for logging."""
    if size_bytes < 100_000:       return "<100KB"
    if size_bytes < 1_000_000:     return "<1MB"
    if size_bytes < 5_000_000:     return "<5MB"
    if size_bytes < 10_000_000:    return "<10MB"
    return ">=10MB"

def _safe_error_category(exc: Exception) -> str:
    """Return generic error category — never raw exception text."""
    if isinstance(exc, asyncio.TimeoutError):   return "timeout"
    if isinstance(exc, OversizedFileError):     return "oversized"
    if isinstance(exc, CorruptFileError):       return "corrupt"
    if isinstance(exc, UnsupportedFileError):   return "unsupported"
    if isinstance(exc, ParserError):            return "parser_error"
    if isinstance(exc, ExcelReadError):         return "read_error"
    return "unexpected_error"

# ── Health & observability ────────────────────────────────────────────────────
@app.get("/health")
async def health():
    return {"status": "ok"}

@app.get("/metrics")
async def prometheus_metrics(request: Request):
    client_ip = (
        request.headers.get("X-Forwarded-For", "").split(",")[0].strip()
        or (request.client.host if request.client else "")
    )
    if METRICS_TOKEN:
        auth = request.headers.get("Authorization", "")
        if not secrets.compare_digest(auth, f"Bearer {METRICS_TOKEN}"):
            raise HTTPException(status_code=403, detail="Forbidden")
    elif client_ip not in METRICS_ALLOWED_IPS:
        raise HTTPException(status_code=403, detail="Forbidden")
    body, content_type = metrics_response()
    return Response(content=body, media_type=content_type)

# ── Routes ────────────────────────────────────────────────────────────────────
@app.post("/api/upload")
async def upload_excel(request: Request, file: UploadFile = File(...)):
    """Receive an Excel file, detect the LEI column, return a job_id + preview."""
    allowed_exts = {".xlsx", ".xlsm", ".xls", ".ods"}
    suffix = Path(file.filename).suffix.lower()
    if suffix not in allowed_exts:
        raise HTTPException(
            status_code=400,
            detail=f"Unsupported file type '{suffix}'. Use .xlsx, .ods or .xls",
        )

    safe_filename = _sanitize_filename(file.filename)

    client_ip = (
        request.headers.get("X-Forwarded-For", "").split(",")[0].strip()
        or (request.client.host if request.client else "unknown")
    )
    if not _check_upload_rate_limit(client_ip):
        raise HTTPException(status_code=429, detail="Too many uploads. Please wait before trying again.")
    if _count_active_jobs_global() >= MAX_JOBS_GLOBAL:
        raise HTTPException(status_code=503, detail="Server busy. Try again later.")
    if _count_active_jobs_for_ip(client_ip) >= MAX_JOBS_PER_IP:
        raise HTTPException(status_code=429, detail="Too many active jobs. Wait for yours to complete.")

    content_length = file.size
    if content_length is not None and content_length > MAX_UPLOAD_BYTES:
        raise HTTPException(status_code=413, detail="File too large (max 10 MB)")

    content = await file.read()

    if not _validate_file_magic(content, suffix):
        raise HTTPException(status_code=400, detail="File content does not match the declared file type.")
    if not validate_modern_excel_archive(content, suffix):
        raise HTTPException(status_code=400, detail="Invalid Excel file structure.")

    bomb_result = _check_decompression_bomb(content, suffix)
    if bomb_result == "too_large":
        raise HTTPException(status_code=413, detail="Archive too large when decompressed.")
    elif bomb_result is not None:
        raise HTTPException(status_code=400, detail="Suspicious archive structure rejected.")

    if suffix == ".xls":
        if not _validate_xls_ole2(content):
            raise HTTPException(status_code=400, detail="Invalid .xls file. Legacy format required.")
        logger.warning("Legacy .xls upload accepted", extra={"ip": client_ip, "filename": safe_filename})

    if len(content) > MAX_UPLOAD_BYTES:
        raise HTTPException(status_code=413, detail="File too large (max 10 MB)")

    loop = asyncio.get_event_loop()
    try:
        leis, column_info = await asyncio.wait_for(
            loop.run_in_executor(_parse_executor, read_lei_from_excel, content, safe_filename),
            timeout=PARSE_TIMEOUT_SECONDS
        )
    except asyncio.TimeoutError:
        logger.error("parse timeout", extra={"ip": client_ip, "filename": safe_filename})
        raise HTTPException(status_code=408, detail="File parsing timed out. Try a smaller file.")
    except OversizedFileError:
        logger.warning("oversized file rejected", extra={"ip": client_ip}, exc_info=True)
        raise HTTPException(status_code=413, detail="File exceeds allowed size or complexity limits.")
    except UnsupportedFileError:
        logger.warning("unsupported file type", extra={"ip": client_ip}, exc_info=True)
        raise HTTPException(status_code=400, detail="Unsupported file type.")
    except CorruptFileError:
        logger.warning("corrupt file rejected", extra={"ip": client_ip}, exc_info=True)
        raise HTTPException(status_code=400, detail="File appears to be corrupt or malformed.")
    except ParserError:
        logger.error("parser failure", extra={"ip": client_ip}, exc_info=True)
        raise HTTPException(status_code=422, detail="Could not parse the uploaded file.")
    except ExcelReadError:
        logger.warning("excel read error", extra={"ip": client_ip}, exc_info=True)
        raise HTTPException(status_code=422, detail="Could not read the file. Please check the format.")
    except Exception:
        logger.error("unexpected upload error", extra={"ip": client_ip}, exc_info=True)
        raise HTTPException(status_code=500, detail="An unexpected error occurred.")

    if not leis:
        raise HTTPException(status_code=422, detail="No LEI codes found in the file")

    if len(leis) > MAX_LEIS:
        raise HTTPException(
            status_code=400,
            detail=f"Too many LEI codes ({len(leis)}). Maximum allowed is {MAX_LEIS}."
        )

    job_id = str(uuid.uuid4())
    jobs[job_id] = {
        "status":         "pending",
        "leis":           leis,
        "results":        [],
        "progress":       0,
        "error_msg":      "",
        "original_bytes": content,
        "filename":       safe_filename,
        "column_info":    column_info,
        "client_ip":      client_ip,
        "created_at":     time.time(),
    }

    non_blank = [l for l in leis if l.strip()]
    logger.info(
    "job_created",
    extra={
        "job_id":      job_id,
        "lei_count":   len(non_blank),
        "file_type":   suffix,
        "size_bucket": _size_bucket(len(content)),
        "ip":          client_ip,
    }
)

    job_token = _make_job_token(job_id)
    return {
        "job_id":    job_id,
        "job_token": job_token,
        "lei_count": len(non_blank),
        "filename":  safe_filename,
        "preview":   non_blank[:8],
    }


@app.post("/api/process/{job_id}")
async def start_processing(job_id: str, background_tasks: BackgroundTasks, job_token: Optional[str] = None):
    """Kick off the background LEI-checking task."""
    _verify_ownership(job_id, job_token)
    job = _get_job(job_id)
    if job["status"] not in ("pending",):
        raise HTTPException(status_code=409, detail=f"Job is already {job['status']}")
    job["status"] = "processing"
    background_tasks.add_task(_run_job, job_id)
    return {"status": "processing", "job_id": job_id}


@app.get("/api/status/{job_id}")
async def get_status(job_id: str, job_token: Optional[str] = None):
    """Poll-based status endpoint."""
    _verify_ownership(job_id, job_token)
    job = _get_job(job_id)
    non_blank_total = len([l for l in job["leis"] if l.strip()])
    return {
        "status":    job["status"],
        "progress":  job["progress"],
        "total":     non_blank_total,
        "results":   job["results"],
        "error_msg": job.get("error_msg", ""),
    }


@app.get("/api/stream/{job_id}")
async def stream_progress(job_id: str, job_token: Optional[str] = None):
    """Server-Sent Events stream."""
    _verify_ownership(job_id, job_token)
    if job_id not in jobs:
        raise HTTPException(status_code=404, detail="Job not found")

    async def generator():
        while True:
            job = jobs.get(job_id)
            if job is None:
                break
            non_blank_total = len([l for l in job["leis"] if l.strip()])
            payload = {
                "status":        job["status"],
                "progress":      job["progress"],
                "total":         non_blank_total,
                "latest_result": job["results"][-1] if job["results"] else None,
                "error_msg":     job.get("error_msg", ""),
            }
            yield f"data: {json.dumps(payload)}\n\n"
            if job["status"] in ("completed", "error"):
                break
            await asyncio.sleep(0.6)

    return StreamingResponse(
        generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control":     "no-cache",
            "X-Accel-Buffering": "no",
        },
    )


@app.get("/api/download/{job_id}")
async def download_results(job_id: str, job_token: Optional[str] = None):
    """Return the enriched Excel file."""
    _verify_ownership(job_id, job_token)
    job = _get_job(job_id)
    if job["status"] != "completed":
        raise HTTPException(status_code=409, detail="Job not completed yet")
    if not job["results"]:
        raise HTTPException(status_code=422, detail="No results to export")

    try:
        out_bytes = write_results_to_excel(
            job["original_bytes"],
            job["results"],
            job["column_info"],
        )
    except Exception:
        logger.exception("Excel write error for job %s", job_id)
        raise HTTPException(status_code=500, detail="Could not generate the output file.")

    stem    = Path(job["filename"]).stem
    ext     = Path(job["filename"]).suffix or ".xlsx"
    dl_name = _sanitize_filename(f"{stem}_LEI_results{ext}")

    return Response(
        content=out_bytes,
        media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        headers={"Content-Disposition": f'attachment; filename="{dl_name}"'},
    )


# ── Background task ───────────────────────────────────────────────────────────
async def _run_job(job_id: str):
    job = jobs[job_id]
    active_jobs.inc()
    t0 = time.perf_counter()
    try:
        def on_progress(idx: int, result: dict):
            job["progress"] = idx + 1
            job["results"].append(result)
            src = (result.get("source") or "not_found").lower().replace(" + ", "_and_").replace("-", "_")
            leis_checked_total.labels(source=src).inc()

        await asyncio.wait_for(
            check_lei_batch(job["leis"], on_progress=on_progress),
            timeout=PROCESS_TIMEOUT_SECONDS
        )
        job["status"] = "completed"
        jobs_total.labels(status="completed").inc()
        job_duration_seconds.observe(time.perf_counter() - t0)
        logger.info("job_completed", extra={"job_id": job_id, "count": len(job["results"])})

    except asyncio.TimeoutError:
        logger.error("job_timeout", extra={"job_id": job_id})
        job["status"]         = "error"
        job["error_msg"]      = "Processing timed out."
        job["original_bytes"] = b""
        jobs_total.labels(status="error").inc()

    except Exception as exc:
        logger.exception("job_failed", extra={"job_id": job_id, "error_category": _safe_error_category(exc)})
        job["status"]         = "error"
        job["error_msg"]      = "Processing failed. Please try again."
        job["original_bytes"] = b""
        jobs_total.labels(status="error").inc()
    finally:
        active_jobs.dec()


# ── Static frontend ───────────────────────────────────────────────────────────
STATIC_DIR = Path(__file__).parent / "static"

@app.get("/", response_class=HTMLResponse)
async def serve_index():
    html_path = STATIC_DIR / "index.html"
    return HTMLResponse(html_path.read_text(encoding="utf-8"))

if STATIC_DIR.exists():
    app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")


# ── Entrypoint ────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=False)