app.py

import os
import io
import json
import sqlite3
import asyncio
import logging
import logging.handlers
import tempfile
import secrets
from datetime import datetime
from threading import Lock
from contextlib import asynccontextmanager

import fitz
import ollama
import easyocr
import numpy as np
import pandas as pd
from PIL import Image
from dotenv import load_dotenv
from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Header, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.concurrency import run_in_threadpool
from json_repair import repair_json
from openpyxl import Workbook, load_workbook

# ===================================================================
# CONFIG — loaded from a .env file (same folder as this script) so you
# can change output paths / model / auth / CORS without touching code.
# Create a file named ".env" next to app.py, for example:
#
#   OUTPUT_DIR=D:\MetaVal_Output
#   OLLAMA_MODEL=gemma3:4b
#   OLLAMA_TIMEOUT_SECONDS=120
#   OCR_GPU=true
#   ALLOWED_ORIGINS=http://192.168.1.50:3000,http://localhost:3000
#   API_KEY=some-long-random-shared-secret
#   PORT=8000
#
# If .env is missing, safe defaults below are used (dev-friendly).
# ===================================================================
load_dotenv()

OUTPUT_DIR = os.getenv("OUTPUT_DIR", r"D:\MetaVal_Output")
LOG_DIR = os.path.join(OUTPUT_DIR, "logs")
OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "gemma3:4b")
OLLAMA_TIMEOUT_SECONDS = int(os.getenv("OLLAMA_TIMEOUT_SECONDS", "120"))
OCR_GPU = os.getenv("OCR_GPU", "true").lower() in ("1", "true", "yes")
ALLOWED_EXTENSIONS = {".pdf"}
# NOTE: no upload size limit — large datasheets are expected and
# accepted as-is. Only extension/emptiness/PDF-header are validated.

# Comma-separated list of allowed origins. Defaults to "*" (open) only
# if nothing is configured — set ALLOWED_ORIGINS in .env for production.
_origins_raw = os.getenv("ALLOWED_ORIGINS", "").strip()
ALLOWED_ORIGINS = [o.strip() for o in _origins_raw.split(",") if o.strip()] or ["*"]

# Shared-secret API key. If unset, auth is DISABLED (dev mode) — a
# warning is logged at startup so it's impossible to miss.
API_KEY = os.getenv("API_KEY", "").strip()

PORT = int(os.getenv("PORT", "8000"))

os.makedirs(OUTPUT_DIR, exist_ok=True)
os.makedirs(LOG_DIR, exist_ok=True)

# ===================================================================
# LOGGING SETUP
# Replaces scattered print() calls with proper structured logging:
# timestamps, log levels, console output AND a rotating log file so
# you can debug issues after the fact instead of losing everything
# when the console window closes / server restarts.
# Rotation keeps disk cleanup automatic: once machine2.log hits 5 MB,
# it rolls over to machine2.log.1, .2, etc. — max 5 old files kept,
# older ones auto-deleted.
# ===================================================================
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)-8s | %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    handlers=[
        logging.StreamHandler(),  # console
        logging.handlers.RotatingFileHandler(
            os.path.join(LOG_DIR, "machine2.log"),
            maxBytes=5 * 1024 * 1024,  # 5 MB per file
            backupCount=5,             # keep last 5 log files
            encoding="utf-8",
        ),
    ],
)

log = logging.getLogger("machine2")

if not API_KEY:
    log.warning(
        "API_KEY is not set in .env — /extract is UNAUTHENTICATED. "
        "Set API_KEY in .env before exposing this server beyond localhost."
    )
if ALLOWED_ORIGINS == ["*"]:
    log.warning("ALLOWED_ORIGINS not set in .env — CORS is wide open ('*').")

# ===================================================================
# OCR READER — declared here, actually loaded inside lifespan startup
# (keeps module import fast; heavy GPU load happens on server start)
# ===================================================================
reader = None

# ===================================================================
# OLLAMA CLIENT WITH TIMEOUT
# The default `ollama.chat()` has NO timeout — if the model ever hangs
# (GPU contention, OOM, stuck generation) the request blocks forever.
# Using an explicit Client with a timeout means a stuck page fails
# fast and gets skipped (see per-page error isolation below) instead
# of hanging the whole PDF — and the whole gpu_semaphore queue behind it.
# ===================================================================
ollama_client = ollama.Client(timeout=OLLAMA_TIMEOUT_SECONDS)

# ===================================================================
# JOB STATUS TRACKING (SQLite)
# Lets Machine 1 poll GET /jobs/{jobId} to recover job status/results
# if it crashes or loses the original /extract response — instead of
# that being the only record a job ever happened.
# Uses stdlib sqlite3, so no extra dependency needed.
# ===================================================================
JOB_DB_PATH = os.path.join(OUTPUT_DIR, "jobs.db")
job_db_lock = Lock()


def init_job_db():
    with job_db_lock:
        conn = sqlite3.connect(JOB_DB_PATH)
        try:
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS jobs (
                    job_id            TEXT PRIMARY KEY,
                    original_filename TEXT,
                    owner_mailbox     TEXT,
                    status            TEXT NOT NULL,
                    rows_extracted    INTEGER,
                    excel_path        TEXT,
                    error_detail      TEXT,
                    created_at        TEXT NOT NULL,
                    updated_at        TEXT NOT NULL
                )
                """
            )
            conn.commit()
        finally:
            conn.close()


def _job_upsert(job_id, **fields):
    """Insert a job row if it doesn't exist, or update the given fields
    if it does. `fields` may include: original_filename, owner_mailbox,
    status, rows_extracted, excel_path, error_detail."""
    if not job_id:
        return
    now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    with job_db_lock:
        conn = sqlite3.connect(JOB_DB_PATH)
        try:
            existing = conn.execute(
                "SELECT job_id FROM jobs WHERE job_id = ?", (job_id,)
            ).fetchone()
            if existing:
                set_clause = ", ".join(f"{k} = ?" for k in fields)
                conn.execute(
                    f"UPDATE jobs SET {set_clause}, updated_at = ? WHERE job_id = ?",
                    (*fields.values(), now, job_id),
                )
            else:
                columns = ["job_id", "created_at", "updated_at"] + list(fields.keys())
                placeholders = ", ".join("?" for _ in columns)
                values = [job_id, now, now] + list(fields.values())
                conn.execute(
                    f"INSERT INTO jobs ({', '.join(columns)}) VALUES ({placeholders})",
                    values,
                )
            conn.commit()
        finally:
            conn.close()


def _job_get(job_id):
    with job_db_lock:
        conn = sqlite3.connect(JOB_DB_PATH)
        conn.row_factory = sqlite3.Row
        try:
            row = conn.execute("SELECT * FROM jobs WHERE job_id = ?", (job_id,)).fetchone()
            return dict(row) if row else None
        finally:
            conn.close()


# ===================================================================
# API KEY AUTH
# Simple shared-secret header check. Machine 1 must send:
#   X-API-Key: <same value as API_KEY in .env>
# If API_KEY is unset in .env, auth is skipped entirely (dev mode —
# a warning was already logged at startup so this can't go unnoticed).
# ===================================================================
async def verify_api_key(x_api_key: str = Header(default=None)):
    if not API_KEY:
        return  # auth disabled (no API_KEY configured)
    if not x_api_key or not secrets.compare_digest(x_api_key, API_KEY):
        raise HTTPException(status_code=401, detail="Invalid or missing X-API-Key header")


# ===================================================================
# LIFESPAN MUST BE DEFINED *BEFORE* THE APP IS CREATED
# ===================================================================
@asynccontextmanager
async def lifespan(app: FastAPI):
    global reader

    # Startup
    log.info("=" * 55)
    log.info("METAVAL MACHINE 2 — STARTING")
    log.info("=" * 55)
    log.info("FastAPI             : Running")

    log.info("Loading OCR Engine (gpu=%s)...", OCR_GPU)
    reader = easyocr.Reader(["en"], gpu=OCR_GPU)
    log.info("OCR Engine          : Loaded")

    init_job_db()
    log.info("Job Tracking DB     : Ready (%s)", JOB_DB_PATH)

    log.info("Ollama              : Ready (model=%s, timeout=%ss)", OLLAMA_MODEL, OLLAMA_TIMEOUT_SECONDS)
    log.info("PDF Engine          : Ready")
    log.info("Excel Engine        : Ready")
    log.info("Excel Output Folder : %s", OUTPUT_DIR)
    log.info("API Endpoint        : http://0.0.0.0:%s", PORT)
    log.info("=" * 55)

    yield  # Application runs here

    # Shutdown
    log.info("Cleaning resources...")
    reader = None
    log.info("MACHINE 2 STOPPED")


# ===================================================================
# APP IS CREATED EXACTLY ONCE, WITH LIFESPAN WIRED IN FROM THE START
# ===================================================================
app = FastAPI(title="MetaVal Machine 2", lifespan=lifespan)

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

COMMON_EXCEL_PATH = os.path.join(OUTPUT_DIR, "Extracted_Data.xlsx")

excel_lock = Lock()             # THREAD LOCK

# ===================================================================
# GPU/LLM WORK QUEUE
# EasyOCR + Ollama share ONE GPU / one model instance — running two
# PDFs' extraction at the same time can crash (CUDA OOM) or corrupt
# output. This semaphore lets requests queue safely: the event loop
# stays free (so /status, /data etc. still respond instantly), but
# only ONE PDF's actual OCR+LLM work runs at a time.
# Raise the value ONLY if your GPU has enough VRAM headroom to truly
# run more than one extraction concurrently.
# ===================================================================
gpu_semaphore = asyncio.Semaphore(1)

# FIXED SYSTEM PROMPT
SYSTEM_PROMPT = """
You are an Industrial Engineering Document Extraction Engine.
Your job is:
1. Read complete document.
2. Extract every useful engineering value.
3. Never explain.
4. Never generate markdown.
5. Return ONLY valid JSON.
Format:
{
    "data":[
        {
            "Field":"Value"
        }
    ]
}
"""


# SAFE PDF EXCEL NAME
def safe_pdf_excel_name(filename):
    filename = os.path.splitext(filename)[0]
    filename = "".join(
        c if c.isalnum() or c in (" ", "_", "-")
        else "_"
        for c in filename
    )
    return filename + ".xlsx"


# CREATE PDF EXCEL PATH
def get_pdf_excel_path(pdf_name):
    return os.path.join(OUTPUT_DIR, safe_pdf_excel_name(pdf_name))


def _atomic_save_workbook(wb, target_path):
    """
    Save an openpyxl Workbook atomically: write to a temp file in the
    SAME directory, then os.replace() it over the target. os.replace
    is atomic on both Windows and POSIX, so a crash or power loss
    mid-write can never leave a half-written / corrupted .xlsx behind
    — the target file is either the old version or the fully new one,
    never something in between.
    """
    target_dir = os.path.dirname(target_path) or "."
    fd, tmp_path = tempfile.mkstemp(suffix=".xlsx", prefix="~tmp_", dir=target_dir)
    os.close(fd)
    try:
        wb.save(tmp_path)
        os.replace(tmp_path, target_path)
    except Exception:
        # Clean up the temp file if something went wrong before the swap
        if os.path.exists(tmp_path):
            os.remove(tmp_path)
        raise


# SAVE ROWS TO COMMON EXCEL
def save_common_excel(rows):
    if not rows:
        return
    with excel_lock:
        if os.path.exists(COMMON_EXCEL_PATH):
            wb = load_workbook(COMMON_EXCEL_PATH)
        else:
            wb = Workbook()
        ws = wb.active
        ws.title = "Output"
        if ws.max_row == 1 and ws["A1"].value is None:
            headers = ["Timestamp"] + list(rows[0].keys())
            ws.append(headers)
        headers = [cell.value for cell in ws[1]]
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        for row in rows:
            ws.append(
                [timestamp] +
                [row.get(h, "") for h in headers[1:]]
            )
        _atomic_save_workbook(wb, COMMON_EXCEL_PATH)


# SAVE INDIVIDUAL PDF EXCEL
def save_pdf_excel(rows, pdf_name):
    if not rows:
        return None
    pdf_excel_path = get_pdf_excel_path(pdf_name)
    df = pd.DataFrame(rows)

    # pandas' df.to_excel doesn't give us atomic-replace directly, so
    # write to a temp file first and swap it in the same way.
    #
    # IMPORTANT: the temp file's name must END in ".xlsx" — pandas
    # picks its writer engine from the file's extension, and it looks
    # at the LAST suffix only. A name like "foo.xlsx.tmp" is seen as
    # extension ".tmp", which has no registered engine and crashes
    # with "ValueError: No engine for filetype: 'tmp'". We also pass
    # engine="openpyxl" explicitly so this can never happen again
    # regardless of the filename.
    target_dir = os.path.dirname(pdf_excel_path) or "."
    fd, tmp_path = tempfile.mkstemp(suffix=".xlsx", prefix="~tmp_", dir=target_dir)
    os.close(fd)
    try:
        df.to_excel(tmp_path, index=False, engine="openpyxl")
        os.replace(tmp_path, pdf_excel_path)
    except Exception:
        if os.path.exists(tmp_path):
            os.remove(tmp_path)
        raise
    return pdf_excel_path


log.info("Phase 1 Loaded Successfully")


# PHASE 2 : PDF + OCR + LLM ENGINE
def extract_text_from_page(page, page_number):
    text = page.get_text().strip()
    if text:
        log.debug("Page %s : Selectable Text Found", page_number)
        return text
    if reader is None:
        raise RuntimeError("OCR engine not ready yet — server still starting up")
    log.debug("Page %s : Running OCR...", page_number)
    pix = page.get_pixmap(dpi=200)
    image = Image.open(io.BytesIO(pix.tobytes("png")))
    text = " ".join(reader.readtext(np.array(image), detail=0))
    return text


# LLM CALL
def call_llm(text):
    log.debug("Sending page to LLM")
    try:
        response = ollama_client.chat(
            model=OLLAMA_MODEL,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": text},
            ],
            options={"temperature": 0},
        )
    except Exception as e:
        # Covers timeout errors and any other Ollama/connection failure.
        raise RuntimeError(f"Ollama call failed or timed out: {e}") from e

    content = response["message"]["content"]

    # Remove Markdown
    if content.startswith("```json"):
        content = content[7:]
    if content.startswith("```"):
        content = content[3:]
    if content.endswith("```"):
        content = content[:-3]

    # Repair JSON
    repaired = repair_json(content)
    data = json.loads(repaired)
    rows = data.get("data", [])
    if isinstance(rows, dict):
        rows = [rows]
    return rows


# PROCESS COMPLETE PDF
def process_pdf(contents, pdf_name):
    results = []
    doc = fitz.open(stream=contents, filetype="pdf")
    try:
        total_pages = len(doc)
        log.info("PDF: %s | Pages: %s", pdf_name, total_pages)

        for page_index, page in enumerate(doc):
            page_number = page_index + 1
            log.info("Processing page %s/%s of %s", page_number, total_pages, pdf_name)
            try:
                text = extract_text_from_page(page, page_number)
                if not text:
                    log.warning("Page %s of %s is empty — skipping", page_number, pdf_name)
                    continue
                rows = call_llm(text)
                for row in rows:
                    row["PDF_Name"] = pdf_name
                    row["Page_No"] = page_number
                    row["Extracted_Time"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                results.extend(rows)
            except Exception as page_err:
                # One bad page (OCR glitch, malformed LLM JSON, Ollama
                # timeout, etc.) should not lose every other page's
                # already-extracted data.
                log.error("Page %s of %s FAILED: %s", page_number, pdf_name, page_err)
                continue

        log.info("Extraction finished for %s | Rows: %s", pdf_name, len(results))
        return results
    finally:
        doc.close()


def _validate_upload(file: UploadFile, contents: bytes):
    """Reject obviously-bad uploads early with a clear 4xx error
    instead of letting them crash deep inside PDF/OCR processing.
    NOTE: no file-size limit here by design — large datasheets are
    expected. Extension, emptiness, and PDF-header checks still apply."""
    ext = os.path.splitext(file.filename or "")[1].lower()
    if ext not in ALLOWED_EXTENSIONS:
        raise HTTPException(
            status_code=400,
            detail=f"Unsupported file type '{ext}'. Only PDF files are accepted.",
        )

    if not contents:
        raise HTTPException(status_code=400, detail="Uploaded file is empty.")

    # Cheap magic-byte check: real PDFs start with %PDF-
    if not contents.startswith(b"%PDF-"):
        raise HTTPException(
            status_code=400,
            detail="File does not look like a valid PDF (bad header).",
        )


# PHASE 3 : FASTAPI ENDPOINT
@app.post("/extract", dependencies=[Depends(verify_api_key)])
async def extract_data(
    file: UploadFile = File(...),
    metadata: str = Form(...),
):
    log.info("=" * 55)
    log.info("MACHINE 2 REQUEST RECEIVED — PDF: %s", file.filename)

    # Read Metadata Sent By Machine 1
    try:
        document_job = json.loads(metadata)
    except Exception:
        raise HTTPException(status_code=400, detail="Invalid Metadata")

    job_id = document_job.get("jobId")
    log.info("Job ID: %s | Mailbox: %s", job_id, document_job.get("ownerMailbox"))

    # Record the job as RECEIVED immediately, before any heavy work —
    # so Machine 1 can poll GET /jobs/{jobId} and see it exists even
    # while it's still queued behind the GPU semaphore.
    await run_in_threadpool(
        _job_upsert,
        job_id,
        original_filename=document_job.get("originalFileName") or file.filename,
        owner_mailbox=document_job.get("ownerMailbox"),
        status="RECEIVED",
        rows_extracted=None,
        excel_path=None,
        error_detail=None,
    )

    # Read + validate PDF
    contents = await file.read()
    try:
        _validate_upload(file, contents)
    except HTTPException as e:
        await run_in_threadpool(_job_upsert, job_id, status="FAILED", error_detail=e.detail)
        raise

    try:
        # Run Complete Extraction — offloaded to a worker thread so the
        # event loop stays free (health checks / other endpoints keep
        # responding), and queued behind gpu_semaphore so two PDFs
        # never hit the GPU / Ollama model at the exact same time.
        async with gpu_semaphore:
            log.info("[%s] Acquired GPU slot — processing now", file.filename)
            await run_in_threadpool(_job_upsert, job_id, status="PROCESSING")
            results = await run_in_threadpool(process_pdf, contents, file.filename)

        # Save Common Excel (also offloaded — openpyxl/pandas I/O is blocking too)
        await run_in_threadpool(save_common_excel, results)

        # Save Individual PDF Excel
        pdf_excel_path = await run_in_threadpool(save_pdf_excel, results, file.filename)
        log.info("Excel saved: %s", pdf_excel_path)

        await run_in_threadpool(
            _job_upsert,
            job_id,
            status="PROCESSED",
            rows_extracted=len(results),
            excel_path=pdf_excel_path,
        )

        # Response Back To Machine 1
        return {
            "success": True,
            "rows": results,
            "excelPath": pdf_excel_path,
            "metadata": {
                "jobId": job_id,
                "originalFileName": document_job.get("originalFileName"),
                "ownerMailbox": document_job.get("ownerMailbox"),
                "processingStatus": "PROCESSED",
                "rowsExtracted": len(results),
                "processedTime": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            },
        }

    except HTTPException as e:
        await run_in_threadpool(_job_upsert, job_id, status="FAILED", error_detail=str(e.detail))
        raise
    except Exception as e:
        log.exception("MACHINE 2 ERROR while processing %s", file.filename)
        await run_in_threadpool(_job_upsert, job_id, status="FAILED", error_detail=str(e))
        raise HTTPException(status_code=500, detail=str(e))


# GET /jobs/{jobId}  -> Poll job status/result if the original /extract
# response was lost (Machine 1 crash, network drop, etc.)
@app.get("/jobs/{job_id}", dependencies=[Depends(verify_api_key)])
async def get_job_status(job_id: str):
    job = await run_in_threadpool(_job_get, job_id)
    if not job:
        raise HTTPException(status_code=404, detail="No job found with this jobId")
    return {"success": True, "job": job}


# HEALTH CHECK
@app.get("/")
async def root():
    return {
        "success": True,
        "message": "MetaVal Machine 2 Running",
        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
    }


# STATUS API
@app.get("/status")
async def status():
    return {
        "machine": "Machine 2",
        "engine": "PDF Extraction Engine",
        "ocr": "EasyOCR",
        "llm": "Gemma3 4B",
        "excelDirectory": OUTPUT_DIR,
        "ocrReady": reader is not None,
        "status": "ONLINE",
    }


# GET /data  -> Master Excel (Extracted_Data.xlsx) ka poora data JSON mein
@app.get("/data", dependencies=[Depends(verify_api_key)])
async def get_master_data():
    if not os.path.exists(COMMON_EXCEL_PATH):
        return {"success": True, "rows": []}
    df = await run_in_threadpool(pd.read_excel, COMMON_EXCEL_PATH)
    df = df.fillna("")
    return {
        "success": True,
        "count": len(df),
        "rows": df.to_dict(orient="records"),
    }


# GET /data/{pdf_name}  -> Ek specific PDF ki Excel ka data JSON mein
@app.get("/data/{pdf_name}", dependencies=[Depends(verify_api_key)])
async def get_pdf_data(pdf_name: str):
    pdf_excel_path = get_pdf_excel_path(pdf_name)
    if not os.path.exists(pdf_excel_path):
        raise HTTPException(status_code=404, detail="File not found for this PDF")
    df = await run_in_threadpool(pd.read_excel, pdf_excel_path)
    df = df.fillna("")
    return {
        "success": True,
        "pdf_name": pdf_name,
        "count": len(df),
        "rows": df.to_dict(orient="records"),
    }


# MAIN
if __name__ == "__main__":
    import uvicorn

    uvicorn.run(
        "app:app",
        host="0.0.0.0",
        port=PORT,
        reload=True,
        reload_includes=["app.py"],
    )


# http://localhost:8000/docs
# uvicorn app:app --host 0.0.0.0 --port 8000
# uvicorn app:app --host 0.0.0.0 --port 8000 --reload

# curl http://localhost:8000/
# curl http://localhost:8000/status

"""
curl -X POST http://localhost:8000/extract ^
-F "file=@D:\Test\EI00362-DATA-OP-1.pdf" ^
-F "metadata={\"jobId\":\"test123\",\"ownerMailbox\":\"test@metaval.com\",\"originalFileName\":\"D:\Test\EI00362-DATA-OP-1.pdf\"}"
"""


"""
Complete API Endpoint URLs
Method      Endpoint URL	                        Use
GET	        http://localhost:8000/	                Health check
GET	        http://localhost:8000/status	        Server status
POST	    http://localhost:8000/extract	        PDF upload + extraction
GET	        http://localhost:8000/data	            Master Excel ka complete data
GET	        http://localhost:8000/data/{pdf_name}	Specific PDF ka data
GET	        http://localhost:8000/docs	            Swagger API testing
GET	        http://localhost:8000/redoc	            ReDoc API documentation
"""