main.py

import os
import time
import uuid
import atexit
import logging
from flask import Flask, jsonify, request, g
from flask_cors import CORS
from flask_limiter import Limiter
from flask_limiter.util import get_remote_address
from config import Config
from logging_config import setup_logging

logger = logging.getLogger(__name__)

ALLOWED_EXTENSIONS = {"png", "jpg", "jpeg", "bmp", "webp"}
SUPPORTED_DOC_TYPES = {"sa_id_card", "sa_id_book", "passport"}

_start_time = time.time()


def _allowed_file(filename):
    return "." in filename and filename.rsplit(".", 1)[1].lower() in ALLOWED_EXTENSIONS


def create_app(config=None):
    cfg = config or Config

    # Setup logging before anything else
    setup_logging(
        log_level=getattr(cfg, "LOG_LEVEL", "INFO"),
        log_format=getattr(cfg, "LOG_FORMAT", "json"),
    )

    app = Flask(__name__)
    app.config["MAX_CONTENT_LENGTH"] = cfg.MAX_CONTENT_LENGTH

    # CORS — allow all origins
    CORS(app, origins="*")

    # Store config on app for access in routes
    app.cfg = cfg

    cfg.UPLOAD_DIR.mkdir(exist_ok=True)

    # Rate limiting
    limiter = Limiter(
        get_remote_address,
        app=app,
        default_limits=["60 per minute"],
        storage_uri="memory://",
    )

    # ── Auth middleware ─────────────────────────────────────────────────

    @app.before_request
    def authenticate():
        if request.path in ("/health", "/", "/warm") or request.path.startswith("/static/"):
            return None
        api_key = getattr(cfg, "API_KEY", "")
        if not api_key:
            return None
        provided = request.headers.get("X-API-Key", "")
        if not provided or provided != api_key:
            return jsonify({"error": "Unauthorized"}), 401

    # ── Request tracking ───────────────────────────────────────────────

    @app.before_request
    def assign_request_id():
        g.request_id = request.headers.get("X-Request-ID", str(uuid.uuid4()))
        g.start_time = time.time()

    @app.after_request
    def log_request(response):
        duration = round((time.time() - g.get("start_time", time.time())) * 1000, 1)
        logger.info(
            "%s %s %s %sms",
            request.method, request.path, response.status_code, duration,
            extra={"request_id": g.get("request_id", "-")},
        )
        response.headers["X-Request-ID"] = g.get("request_id", "")
        return response

    # ── Error handlers ─────────────────────────────────────────────────

    @app.errorhandler(400)
    def bad_request(e):
        return jsonify({"error": getattr(e, "description", "Bad request")}), 400

    @app.errorhandler(401)
    def unauthorized(e):
        return jsonify({"error": "Unauthorized"}), 401

    @app.errorhandler(404)
    def not_found(e):
        return jsonify({"error": getattr(e, "description", "Not found")}), 404

    @app.errorhandler(413)
    def request_too_large(e):
        return jsonify({"error": "Request body too large"}), 413

    @app.errorhandler(429)
    def rate_limited(e):
        return jsonify({"error": "Rate limit exceeded"}), 429

    @app.errorhandler(500)
    def internal_error(e):
        logger.exception("Internal server error")
        return jsonify({"error": "Internal server error"}), 500

    # ── Routes ─────────────────────────────────────────────────────────

    @app.route("/")
    def dashboard():
        return app.send_static_file("index.html")

    @app.route("/health")
    def health_check():
        upload_ok = cfg.UPLOAD_DIR.exists() and os.access(str(cfg.UPLOAD_DIR), os.W_OK)
        healthy = upload_ok

        return jsonify({
            "status": "healthy" if healthy else "unhealthy",
            "uptime_seconds": round(time.time() - _start_time, 1),
            "checks": {
                "upload_directory": "ok" if upload_ok else "missing or not writable",
            },
        }), 200 if healthy else 503

    @app.route("/warm")
    def warm():
        """Warm up VLM + PaddleOCR + MRZ models. Call before first OCR request."""
        results = {}

        try:
            from engine.vlm_extractor import warm_vlm
            results["vlm"] = "ready" if warm_vlm() else "unavailable"
        except Exception:
            results["vlm"] = "unavailable"

        try:
            from engine.ocr import _get_ocr
            _get_ocr()
            results["paddleocr"] = "ready"
        except Exception:
            results["paddleocr"] = "unavailable"

        try:
            from engine.mrz_reader import warm_mrz
            results["mrz"] = "ready" if warm_mrz() else "unavailable"
        except Exception:
            results["mrz"] = "unavailable"

        return jsonify({"status": "warm", "engines": results}), 200

    @app.route("/api/ocr", methods=["POST"])
    @limiter.limit("10 per minute")
    def process_id():
        """Accept document image(s) and return extracted fields as JSON.

        Form fields:
            front: Front image file (required) — or 'file' for legacy
            back: Back image file (optional, only for sa_id_card)
            doc_type: Document type (required): sa_id_card | sa_id_book | passport
        """
        # Validate doc_type
        doc_type = request.form.get("doc_type", "").strip()
        if not doc_type:
            return jsonify({
                "error": "doc_type is required. Must be one of: sa_id_card, sa_id_book, passport"
            }), 400
        if doc_type not in SUPPORTED_DOC_TYPES:
            return jsonify({
                "error": f"Invalid doc_type '{doc_type}'. Must be one of: sa_id_card, sa_id_book, passport"
            }), 400

        # Get front image (support both 'front' and legacy 'file' key)
        front_file = request.files.get("front") or request.files.get("file")
        if not front_file:
            return jsonify({"error": "No front image uploaded. Send a file with key 'front'."}), 400

        if front_file.filename == "" or not _allowed_file(front_file.filename):
            return jsonify({
                "error": f"Invalid file type. Accepted: {', '.join(ALLOWED_EXTENSIONS)}"
            }), 400

        # Get back image (optional, only for sa_id_card)
        back_file = request.files.get("back")
        if back_file and doc_type != "sa_id_card":
            return jsonify({"error": "Back image is only accepted for doc_type=sa_id_card"}), 400
        if back_file and (back_file.filename == "" or not _allowed_file(back_file.filename)):
            return jsonify({
                "error": f"Invalid back file type. Accepted: {', '.join(ALLOWED_EXTENSIONS)}"
            }), 400

        # Save front image with UUID filename
        front_ext = front_file.filename.rsplit(".", 1)[1].lower()
        front_filename = f"{uuid.uuid4().hex}.{front_ext}"
        front_path = os.path.join(str(cfg.UPLOAD_DIR), front_filename)
        front_file.save(front_path)

        # Save back image (if provided)
        back_path = None
        if back_file:
            back_ext = back_file.filename.rsplit(".", 1)[1].lower()
            back_filename = f"{uuid.uuid4().hex}.{back_ext}"
            back_path = os.path.join(str(cfg.UPLOAD_DIR), back_filename)
            back_file.save(back_path)

        try:
            from engine.ocr import process_id_image
            result = process_id_image(front_path, doc_type, back_path)
            return jsonify(result), 200
        except ValueError as e:
            return jsonify({"error": str(e)}), 400
        except Exception as e:
            logger.exception("OCR processing failed")
            return jsonify({"error": f"OCR processing failed: {e}"}), 500
        finally:
            # Always clean up — never store personal ID images (POPIA)
            if os.path.exists(front_path):
                os.remove(front_path)
            if back_path and os.path.exists(back_path):
                os.remove(back_path)

    @app.route("/api/serial", methods=["POST"])
    @limiter.limit("10 per minute")
    def process_serial():
        """Accept a product image and extract the serial number.

        Form fields:
            image: Image file (required)
        """
        image_file = request.files.get("image")
        if not image_file:
            return jsonify({"error": "No image uploaded. Send a file with key 'image'."}), 400

        if image_file.filename == "" or not _allowed_file(image_file.filename):
            return jsonify({
                "error": f"Invalid file type. Accepted: {', '.join(ALLOWED_EXTENSIONS)}"
            }), 400

        ext = image_file.filename.rsplit(".", 1)[1].lower()
        filename = f"{uuid.uuid4().hex}.{ext}"
        filepath = os.path.join(str(cfg.UPLOAD_DIR), filename)
        image_file.save(filepath)

        try:
            from engine.serial_reader import process_serial_image
            result = process_serial_image(filepath)
            return jsonify(result), 200
        except Exception as e:
            logger.exception("Serial number processing failed")
            return jsonify({"error": f"Serial number processing failed: {e}"}), 500
        finally:
            if os.path.exists(filepath):
                os.remove(filepath)

    @app.route("/api/barcode/parse", methods=["POST"])
    @limiter.limit("60 per minute")
    def parse_barcode_text():
        """Parse raw barcode text (decoded client-side) into structured fields."""
        data = request.get_json()
        if not data or not data.get("text"):
            return jsonify({"error": "No barcode text provided"}), 400

        text = data["text"]

        from engine.barcode_reader import _parse_pdf417, _parse_code39
        from engine.id_parser import validate_luhn

        result = _parse_pdf417(text) or _parse_code39(text)

        if result:
            fields = {
                "id_number": result.get("id_number"),
                "surname": result.get("surname"),
                "names": result.get("names"),
                "sex": result.get("sex"),
                "date_of_birth": result.get("date_of_birth"),
                "nationality": result.get("nationality"),
                "country_of_birth": result.get("country_of_birth"),
                "citizenship_status": result.get("citizenship_status"),
            }
            id_valid = "not_applicable"
            if fields.get("id_number"):
                id_valid = "passed" if validate_luhn(fields["id_number"]) else "failed"

            return jsonify({
                "doc_type": "sa_id_card",
                "fields": fields,
                "barcode_data": result,
                "extraction_method": "barcode",
                "checks": {
                    "barcode_valid": "passed",
                    "id_number_valid": id_valid,
                },
                "overall_result": "pass" if id_valid == "passed" else "fail",
                "confidence": {},
                "raw_text": [],
            }), 200

        return jsonify({
            "doc_type": "sa_id_card",
            "fields": {},
            "barcode_data": None,
            "extraction_method": "none",
            "checks": {"barcode_valid": "failed"},
            "overall_result": "fail",
            "confidence": {},
            "raw_text": [],
        }), 200

    # ── Shutdown logging ───────────────────────────────────────────────

    def on_shutdown():
        try:
            logger.info("ID OCR Engine shutting down")
        except Exception:
            pass

    atexit.register(on_shutdown)

    logger.info("ID OCR Engine started (debug=%s)", cfg.DEBUG)
    return app


if __name__ == "__main__":
    from dotenv import load_dotenv
    load_dotenv()
    # Reload config after dotenv
    import importlib
    import config as config_module
    importlib.reload(config_module)
    from config import Config as ReloadedConfig

    app = create_app(ReloadedConfig)
    app.run(debug=ReloadedConfig.DEBUG, host=ReloadedConfig.HOST, port=ReloadedConfig.PORT)