#!/usr/bin/env python3
"""
Production-ready Docling VLM FastAPI Server with true concurrent processing.

This server uses FastAPI with asyncio to handle multiple document processing
requests concurrently, maximizing GPU utilization and throughput.

Features:
    - True concurrent processing (multiple pages at once)
    - Single model instance (efficient VRAM usage)
    - Thread pool executor for blocking operations
    - Production-ready with proper error handling
    - Health checks and monitoring endpoints

Usage:
    # Install dependencies
    uv pip install "fastapi[standard]"

    # Start server
    python scripts/docling_fastapi_server.py

    # Or with custom settings
    python scripts/docling_fastapi_server.py --port 8900 --workers 4
"""

import argparse
import asyncio
import tempfile
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Annotated, Optional

import uvicorn
from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import VlmPipelineOptions
from docling.document_converter import DocumentConverter, PdfFormatOption
from docling.pipeline.vlm_pipeline import VlmPipeline
from fastapi import FastAPI, File, HTTPException, UploadFile
from fastapi.responses import JSONResponse
from loguru import logger


class DoclingModelServer:
    """Thread-safe server that keeps VLM model loaded and processes concurrently."""

    def __init__(self, device: str = "cuda", max_workers: int = 4):
        self.device = device
        self.max_workers = max_workers
        self.converter: Optional[DocumentConverter] = None
        self.executor: Optional[ThreadPoolExecutor] = None
        self.initialized = False
        self._lock = threading.Lock()
        self._stats = {"requests_processed": 0, "total_processing_time": 0.0, "active_requests": 0}

    def initialize(self):
        """Load the model into memory (expensive operation)."""
        with self._lock:
            if self.initialized:
                logger.info("Model already initialized")
                return

            logger.info("🚀 Initializing Docling VLM model server...")
            logger.info(f"   Device: {self.device}")
            logger.info(f"   Max workers: {self.max_workers}")

            start_time = time.time()

            # Configure VLM pipeline
            pipeline_options = VlmPipelineOptions()

            # Initialize converter (loads model)
            self.converter = DocumentConverter(
                format_options={
                    InputFormat.PDF: PdfFormatOption(
                        pipeline_cls=VlmPipeline,
                        pipeline_options=pipeline_options,
                    )
                }
            )

            # Create thread pool for blocking operations
            self.executor = ThreadPoolExecutor(
                max_workers=self.max_workers, thread_name_prefix="docling-worker"
            )

            init_time = time.time() - start_time
            logger.success(f"✓ Model loaded in {init_time:.2f} seconds")
            logger.success(f"🎯 Server ready with {self.max_workers} worker threads")
            self.initialized = True

    def _process_sync(self, pdf_path: Path, request_id: str) -> str:
        """
        Synchronous processing function (runs in thread pool).
        This is thread-safe because docling's converter can handle concurrent calls.
        """
        if not self.initialized or self.converter is None:
            raise RuntimeError("Server not initialized")

        thread_name = threading.current_thread().name
        logger.info(f"[{request_id}] 📄 Processing: {pdf_path.name} (thread: {thread_name})")

        with self._lock:
            self._stats["active_requests"] += 1

        start_time = time.time()

        try:
            # Convert document (blocking operation, but thread-safe)
            result = self.converter.convert(pdf_path)
            markdown = result.document.export_to_markdown()

            process_time = time.time() - start_time

            with self._lock:
                self._stats["requests_processed"] += 1
                self._stats["total_processing_time"] += process_time
                self._stats["active_requests"] -= 1

            logger.success(f"[{request_id}]    ✓ {pdf_path.name} processed in {process_time:.2f}s")
            return markdown

        except Exception as e:
            with self._lock:
                self._stats["active_requests"] -= 1
            logger.error(f"[{request_id}] Error processing {pdf_path.name}: {e}")
            raise

    async def process_document(self, pdf_path: Path, request_id: str) -> str:
        """Async wrapper that runs blocking operation in thread pool."""
        if not self.initialized:
            raise RuntimeError("Server not initialized")

        # Run blocking operation in thread pool
        loop = asyncio.get_event_loop()
        markdown = await loop.run_in_executor(self.executor, self._process_sync, pdf_path, request_id)

        return markdown

    def get_stats(self) -> dict[str, float]:
        """Get current server statistics."""
        with self._lock:
            avg_time = (
                self._stats["total_processing_time"] / self._stats["requests_processed"]
                if self._stats["requests_processed"] > 0
                else 0
            )
            return {**self._stats, "average_processing_time": avg_time}

    def shutdown(self):
        """Clean shutdown of thread pool."""
        if self.executor:
            logger.info("Shutting down thread pool...")
            self.executor.shutdown(wait=True)
            logger.info("✓ Thread pool shutdown complete")


# Global server instance
_server: Optional[DoclingModelServer] = None
_request_counter = 0
_counter_lock = threading.Lock()


def get_server() -> DoclingModelServer:
    """Get the global server instance."""
    if _server is None:
        raise RuntimeError("Server not initialized")
    return _server


def get_request_id() -> str:
    """Generate unique request ID."""
    global _request_counter
    with _counter_lock:
        _request_counter += 1
        return f"req-{_request_counter:04d}"


# Lifespan context manager for startup/shutdown
@asynccontextmanager
async def lifespan(app: FastAPI):
    """Manage server lifecycle."""
    global _server

    # Startup
    max_workers = app.state.max_workers if hasattr(app.state, "max_workers") else 4
    _server = DoclingModelServer(max_workers=max_workers)
    _server.initialize()

    yield

    # Shutdown
    _server.shutdown()


# FastAPI app with lifespan
app = FastAPI(
    title="Docling VLM Server",
    description="High-performance concurrent document processing with VLM models",
    version="2.0.0",
    lifespan=lifespan,
)


@app.get("/health")
async def health():
    """Health check endpoint."""
    try:
        server = get_server()
        stats = server.get_stats()
        return {
            "status": "healthy",
            "initialized": server.initialized,
            "max_workers": server.max_workers,
            "active_requests": stats["active_requests"],
            "total_requests": stats["requests_processed"],
        }
    except Exception as e:
        return JSONResponse(status_code=503, content={"status": "unhealthy", "error": str(e)})


@app.post("/convert")
async def convert(file: Annotated[UploadFile, File(...)]):
    """
    Convert a single PDF to markdown.

    Supports concurrent processing - multiple requests are handled simultaneously
    using a thread pool, maximizing GPU utilization.
    """
    request_id = get_request_id()

    if not file.filename or not file.filename.endswith(".pdf"):
        raise HTTPException(status_code=400, detail="Only PDF files are supported")

    server = get_server()

    logger.info(f"[{request_id}] 📥 Received: {file.filename}")

    # Save uploaded file to temp location
    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
        content = await file.read()
        tmp.write(content)
        tmp_path = Path(tmp.name)

    try:
        # Process document (async, runs in thread pool)
        markdown = await server.process_document(tmp_path, request_id)

        return {
            "status": "success",
            "filename": file.filename,
            "markdown": markdown,
            "size": len(markdown),
            "request_id": request_id,
        }

    except Exception as e:
        logger.error(f"[{request_id}] Error processing {file.filename}: {e}")
        raise HTTPException(status_code=500, detail=str(e)) from e

    finally:
        # Cleanup temp file
        if tmp_path.exists():
            tmp_path.unlink()


@app.post("/convert_batch")
async def convert_batch(files: Annotated[list[UploadFile], File(...)]):
    """
    Convert multiple PDFs to markdown concurrently.

    All files are processed in parallel up to max_workers limit.
    This endpoint is ideal for batch processing.
    """
    if not files:
        raise HTTPException(status_code=400, detail="No files provided")

    server = get_server()
    batch_id = get_request_id()

    logger.info(f"[{batch_id}] 📦 Batch request with {len(files)} files")

    # Save all files to temp locations
    tmp_paths = []
    for i, file in enumerate(files):
        if not file.filename or not file.filename.endswith(".pdf"):
            raise HTTPException(status_code=400, detail=f"File {file.filename} is not a PDF")

        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
            content = await file.read()
            tmp.write(content)
            tmp_paths.append((Path(tmp.name), file.filename, f"{batch_id}-{i+1}"))

    try:
        # Process all documents concurrently
        tasks = [server.process_document(tmp_path, req_id) for tmp_path, _, req_id in tmp_paths]

        markdowns = await asyncio.gather(*tasks)

        logger.success(f"[{batch_id}] ✓ Batch completed: {len(markdowns)} files")

        return {
            "status": "success",
            "count": len(markdowns),
            "batch_id": batch_id,
            "results": [
                {"filename": filename, "markdown": markdown, "size": len(markdown), "request_id": req_id}
                for (_, filename, req_id), markdown in zip(tmp_paths, markdowns, strict=False)
            ],
        }

    except Exception as e:
        logger.error(f"[{batch_id}] Error processing batch: {e}")
        raise HTTPException(status_code=500, detail=str(e)) from e

    finally:
        # Cleanup temp files
        for tmp_path, _, _ in tmp_paths:
            if tmp_path.exists():
                tmp_path.unlink()


@app.get("/stats")
async def stats():
    """Get detailed server statistics."""
    server = get_server()
    return {
        "initialized": server.initialized,
        "max_workers": server.max_workers,
        "device": server.device,
        **server.get_stats(),
    }


def main():
    parser = argparse.ArgumentParser(description="Docling FastAPI Server v2")
    parser.add_argument("--host", default="0.0.0.0", help="Host to bind to")
    parser.add_argument("--port", type=int, default=8900, help="Port to bind to")
    parser.add_argument(
        "--workers", type=int, default=4, help="Number of worker threads for concurrent processing"
    )
    parser.add_argument("--reload", action="store_true", help="Enable auto-reload for development")

    args = parser.parse_args()

    logger.info("=" * 60)
    logger.info("Docling FastAPI Server v2 - Production Ready")
    logger.info("=" * 60)
    logger.info(f"Host: {args.host}")
    logger.info(f"Port: {args.port}")
    logger.info(f"Workers: {args.workers}")
    logger.info(f"Concurrent requests: Up to {args.workers} simultaneous")
    logger.info("=" * 60)

    # Create app instance with config
    app.state.max_workers = args.workers

    # Run with uvicorn
    uvicorn.run(app, host=args.host, port=args.port, log_level="info")


if __name__ == "__main__":
    main()