backend/api.py

"""
SPARKNET Backend API - GPU-Accelerated Document Processing

This FastAPI service runs on a GPU server (e.g., Lytos) and provides:
- Document processing with PaddleOCR
- Layout detection
- RAG indexing and querying
- Embedding generation
- LLM inference via Ollama

Deploy this on your GPU server and connect Streamlit Cloud to it.
"""

from fastapi import FastAPI, HTTPException, UploadFile, File, Form, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from typing import Optional, List, Dict, Any
import hashlib
import tempfile
import os
import sys
from pathlib import Path
from datetime import datetime
import asyncio

# Add project root to path
PROJECT_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

app = FastAPI(
    title="SPARKNET Backend API",
    description="GPU-accelerated document processing for Technology Transfer Office automation",
    version="1.0.0",
    docs_url="/docs",
    redoc_url="/redoc",
)

# CORS - Allow Streamlit Cloud to connect
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Configure specific origins in production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ============================================================================
# Pydantic Models
# ============================================================================

class HealthResponse(BaseModel):
    status: str
    timestamp: str
    version: str = "1.0.0"


class SystemStatus(BaseModel):
    ollama_available: bool
    ollama_models: List[str] = []
    gpu_available: bool = False
    gpu_name: Optional[str] = None
    rag_ready: bool = False
    indexed_chunks: int = 0
    embedding_model: Optional[str] = None
    llm_model: Optional[str] = None


class ProcessRequest(BaseModel):
    filename: str
    options: Dict[str, Any] = Field(default_factory=dict)


class ProcessResponse(BaseModel):
    success: bool
    doc_id: str
    filename: str
    raw_text: str = ""
    chunks: List[Dict[str, Any]] = []
    page_count: int = 0
    ocr_regions: List[Dict[str, Any]] = []
    layout_regions: List[Dict[str, Any]] = []
    ocr_confidence: float = 0.0
    layout_confidence: float = 0.0
    processing_time: float = 0.0
    error: Optional[str] = None


class IndexRequest(BaseModel):
    doc_id: str
    text: str
    chunks: List[Dict[str, Any]] = []
    metadata: Dict[str, Any] = Field(default_factory=dict)


class IndexResponse(BaseModel):
    success: bool
    doc_id: str
    num_chunks: int = 0
    error: Optional[str] = None


class QueryRequest(BaseModel):
    question: str
    filters: Optional[Dict[str, Any]] = None
    top_k: int = 5


class QueryResponse(BaseModel):
    success: bool
    answer: str = ""
    sources: List[Dict[str, Any]] = []
    confidence: float = 0.0
    latency_ms: float = 0.0
    validated: bool = False
    error: Optional[str] = None


class SearchRequest(BaseModel):
    query: str
    top_k: int = 5
    doc_filter: Optional[str] = None


class DocumentInfo(BaseModel):
    doc_id: str
    filename: str = ""
    chunk_count: int = 0
    indexed_at: Optional[str] = None


# ============================================================================
# Global State
# ============================================================================

_rag_system = None
_processing_queue = {}


def get_rag_system():
    """Initialize and return the RAG system."""
    global _rag_system

    if _rag_system is not None:
        return _rag_system

    try:
        from src.rag.agentic import AgenticRAG, RAGConfig
        from src.rag.store import get_vector_store, VectorStoreConfig, reset_vector_store
        from src.rag.embeddings import get_embedding_adapter, EmbeddingConfig, reset_embedding_adapter

        # Check Ollama
        ollama_ok, models = check_ollama_sync()
        if not ollama_ok:
            return None

        # Select models
        EMBEDDING_MODELS = ["nomic-embed-text", "mxbai-embed-large:latest", "mxbai-embed-large"]
        LLM_MODELS = ["llama3.2:latest", "llama3.1:8b", "mistral:latest", "qwen2.5:14b"]

        embed_model = next((m for m in EMBEDDING_MODELS if m in models), EMBEDDING_MODELS[0])
        llm_model = next((m for m in LLM_MODELS if m in models), LLM_MODELS[0])

        # Reset singletons
        reset_vector_store()
        reset_embedding_adapter()

        # Initialize embedding adapter
        embed_config = EmbeddingConfig(
            ollama_model=embed_model,
            ollama_base_url="http://localhost:11434",
        )
        embedder = get_embedding_adapter(config=embed_config)

        # Initialize vector store
        store_config = VectorStoreConfig(
            persist_directory="data/sparknet_unified_rag",
            collection_name="sparknet_documents",
            similarity_threshold=0.0,
        )
        store = get_vector_store(config=store_config)

        # Initialize RAG config
        rag_config = RAGConfig(
            model=llm_model,
            base_url="http://localhost:11434",
            max_revision_attempts=1,
            enable_query_planning=True,
            enable_reranking=True,
            enable_validation=True,
            retrieval_top_k=10,
            final_top_k=5,
            min_confidence=0.3,
            verbose=False,
        )

        # Initialize RAG system
        rag = AgenticRAG(
            config=rag_config,
            vector_store=store,
            embedding_adapter=embedder,
        )

        _rag_system = {
            "rag": rag,
            "store": store,
            "embedder": embedder,
            "embed_model": embed_model,
            "llm_model": llm_model,
        }

        return _rag_system

    except Exception as e:
        print(f"RAG init error: {e}")
        return None


def check_ollama_sync():
    """Check Ollama availability synchronously."""
    try:
        import httpx
        with httpx.Client(timeout=3.0) as client:
            resp = client.get("http://localhost:11434/api/tags")
            if resp.status_code == 200:
                models = [m["name"] for m in resp.json().get("models", [])]
                return True, models
    except:
        pass
    return False, []


def check_gpu():
    """Check GPU availability."""
    try:
        import torch
        if torch.cuda.is_available():
            return True, torch.cuda.get_device_name(0)
    except:
        pass
    return False, None


# ============================================================================
# API Endpoints
# ============================================================================

@app.get("/", response_model=HealthResponse)
async def root():
    """Health check endpoint."""
    return HealthResponse(
        status="healthy",
        timestamp=datetime.now().isoformat(),
    )


@app.get("/api/health", response_model=HealthResponse)
async def health():
    """Health check endpoint."""
    return HealthResponse(
        status="healthy",
        timestamp=datetime.now().isoformat(),
    )


@app.get("/api/status", response_model=SystemStatus)
async def get_status():
    """Get system status including Ollama, GPU, and RAG availability."""
    ollama_ok, models = check_ollama_sync()
    gpu_ok, gpu_name = check_gpu()

    rag = get_rag_system()
    rag_ready = rag is not None

    indexed_chunks = 0
    embed_model = None
    llm_model = None

    if rag:
        try:
            indexed_chunks = rag["store"].count()
            embed_model = rag.get("embed_model")
            llm_model = rag.get("llm_model")
        except:
            pass

    return SystemStatus(
        ollama_available=ollama_ok,
        ollama_models=models,
        gpu_available=gpu_ok,
        gpu_name=gpu_name,
        rag_ready=rag_ready,
        indexed_chunks=indexed_chunks,
        embedding_model=embed_model,
        llm_model=llm_model,
    )


@app.post("/api/process", response_model=ProcessResponse)
async def process_document(
    file: UploadFile = File(...),
    ocr_engine: str = Form(default="paddleocr"),
    max_pages: int = Form(default=10),
    enable_layout: bool = Form(default=True),
    preserve_tables: bool = Form(default=True),
):
    """
    Process a document with OCR and layout detection.

    This endpoint uses GPU-accelerated PaddleOCR for text extraction.
    """
    import time
    start_time = time.time()

    # Read file
    file_bytes = await file.read()
    filename = file.filename

    # Generate doc ID
    content_hash = hashlib.md5(file_bytes[:1000]).hexdigest()[:8]
    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
    doc_id = hashlib.md5(f"{filename}_{timestamp}_{content_hash}".encode()).hexdigest()[:12]

    # Save to temp file
    suffix = Path(filename).suffix
    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
        tmp.write(file_bytes)
        tmp_path = tmp.name

    try:
        # Try full document processing pipeline
        try:
            from src.document.pipeline.processor import DocumentProcessor, PipelineConfig
            from src.document.ocr import OCRConfig
            from src.document.layout import LayoutConfig
            from src.document.chunking.chunker import ChunkerConfig

            chunker_config = ChunkerConfig(
                preserve_table_structure=preserve_tables,
                detect_table_headers=True,
                chunk_tables=True,
                chunk_figures=True,
                include_captions=True,
            )

            layout_config = LayoutConfig(
                method="rule_based",
                detect_tables=True,
                detect_figures=True,
                detect_headers=True,
                detect_titles=True,
                detect_lists=True,
                min_confidence=0.3,
                heading_font_ratio=1.1,
            )

            config = PipelineConfig(
                ocr=OCRConfig(engine=ocr_engine),
                layout=layout_config,
                chunking=chunker_config,
                max_pages=max_pages,
                include_ocr_regions=True,
                include_layout_regions=enable_layout,
                generate_full_text=True,
            )

            processor = DocumentProcessor(config)
            processor.initialize()
            result = processor.process(tmp_path)

            # Convert to response format
            chunks_list = []
            for chunk in result.chunks:
                chunks_list.append({
                    "chunk_id": chunk.chunk_id,
                    "text": chunk.text,
                    "page": chunk.page,
                    "chunk_type": chunk.chunk_type.value,
                    "confidence": chunk.confidence,
                    "bbox": chunk.bbox.to_xyxy() if chunk.bbox else None,
                })

            ocr_regions = []
            for region in result.ocr_regions:
                ocr_regions.append({
                    "text": region.text,
                    "confidence": region.confidence,
                    "page": region.page,
                    "bbox": region.bbox.to_xyxy() if region.bbox else None,
                })

            layout_regions = []
            for region in result.layout_regions:
                layout_regions.append({
                    "id": region.id,
                    "type": region.type.value,
                    "confidence": region.confidence,
                    "page": region.page,
                    "bbox": region.bbox.to_xyxy() if region.bbox else None,
                })

            processing_time = time.time() - start_time

            return ProcessResponse(
                success=True,
                doc_id=doc_id,
                filename=filename,
                raw_text=result.full_text,
                chunks=chunks_list,
                page_count=result.metadata.num_pages,
                ocr_regions=ocr_regions,
                layout_regions=layout_regions,
                ocr_confidence=result.metadata.ocr_confidence_avg or 0.0,
                layout_confidence=result.metadata.layout_confidence_avg or 0.0,
                processing_time=processing_time,
            )

        except Exception as e:
            # Fallback to simple extraction
            return await process_document_fallback(file_bytes, filename, doc_id, max_pages, str(e), start_time)

    finally:
        # Cleanup
        if os.path.exists(tmp_path):
            os.unlink(tmp_path)


async def process_document_fallback(
    file_bytes: bytes,
    filename: str,
    doc_id: str,
    max_pages: int,
    reason: str,
    start_time: float
) -> ProcessResponse:
    """Fallback document processing using PyMuPDF."""
    import time

    text = ""
    page_count = 1
    suffix = Path(filename).suffix.lower()

    if suffix == ".pdf":
        try:
            import fitz
            import io
            pdf_stream = io.BytesIO(file_bytes)
            doc = fitz.open(stream=pdf_stream, filetype="pdf")
            page_count = len(doc)
            max_p = min(max_pages, page_count)

            text_parts = []
            for page_num in range(max_p):
                page = doc[page_num]
                text_parts.append(f"--- Page {page_num + 1} ---\n{page.get_text()}")
            text = "\n\n".join(text_parts)
            doc.close()
        except Exception as e:
            text = f"PDF extraction failed: {e}"
    elif suffix in [".txt", ".md"]:
        try:
            text = file_bytes.decode("utf-8")
        except:
            text = file_bytes.decode("latin-1", errors="ignore")
    else:
        text = f"Unsupported file type: {suffix}"

    # Simple chunking
    chunk_size = 500
    overlap = 50
    chunks = []

    for i in range(0, len(text), chunk_size - overlap):
        chunk_text = text[i:i + chunk_size]
        if len(chunk_text.strip()) > 20:
            chunks.append({
                "chunk_id": f"{doc_id}_chunk_{len(chunks)}",
                "text": chunk_text,
                "page": 0,
                "chunk_type": "text",
                "confidence": 0.9,
                "bbox": None,
            })

    processing_time = time.time() - start_time

    return ProcessResponse(
        success=True,
        doc_id=doc_id,
        filename=filename,
        raw_text=text,
        chunks=chunks,
        page_count=page_count,
        ocr_regions=[],
        layout_regions=[],
        ocr_confidence=0.9,
        layout_confidence=0.0,
        processing_time=processing_time,
        error=f"Fallback mode: {reason}",
    )


@app.post("/api/index", response_model=IndexResponse)
async def index_document(request: IndexRequest):
    """Index a document into the RAG vector store."""
    rag = get_rag_system()

    if not rag:
        return IndexResponse(
            success=False,
            doc_id=request.doc_id,
            error="RAG system not available. Check Ollama status.",
        )

    try:
        store = rag["store"]
        embedder = rag["embedder"]

        chunk_dicts = []
        embeddings = []

        for i, chunk in enumerate(request.chunks):
            chunk_text = chunk.get("text", "") if isinstance(chunk, dict) else str(chunk)

            if len(chunk_text.strip()) < 20:
                continue

            chunk_id = chunk.get("chunk_id", f"{request.doc_id}_chunk_{i}")
            chunk_dict = {
                "chunk_id": chunk_id,
                "document_id": request.doc_id,
                "text": chunk_text,
                "page": chunk.get("page", 0) if isinstance(chunk, dict) else 0,
                "chunk_type": "text",
                "source_path": request.metadata.get("filename", ""),
                "sequence_index": i,
            }
            chunk_dicts.append(chunk_dict)

            embedding = embedder.embed_text(chunk_text)
            embeddings.append(embedding)

        if not chunk_dicts:
            return IndexResponse(
                success=False,
                doc_id=request.doc_id,
                error="No valid chunks to index",
            )

        store.add_chunks(chunk_dicts, embeddings)

        return IndexResponse(
            success=True,
            doc_id=request.doc_id,
            num_chunks=len(chunk_dicts),
        )

    except Exception as e:
        return IndexResponse(
            success=False,
            doc_id=request.doc_id,
            error=str(e),
        )


@app.post("/api/query", response_model=QueryResponse)
async def query_rag(request: QueryRequest):
    """Query the RAG system."""
    import time
    start_time = time.time()

    rag = get_rag_system()

    if not rag:
        return QueryResponse(
            success=False,
            error="RAG system not available. Check Ollama status.",
        )

    try:
        response = rag["rag"].query(request.question, filters=request.filters)
        latency_ms = (time.time() - start_time) * 1000

        sources = []
        if hasattr(response, 'citations') and response.citations:
            for cite in response.citations:
                sources.append({
                    "index": cite.index if hasattr(cite, 'index') else 0,
                    "text_snippet": cite.text_snippet if hasattr(cite, 'text_snippet') else str(cite),
                    "relevance_score": cite.relevance_score if hasattr(cite, 'relevance_score') else 0.0,
                    "document_id": cite.document_id if hasattr(cite, 'document_id') else "",
                    "page": cite.page if hasattr(cite, 'page') else 0,
                })

        return QueryResponse(
            success=True,
            answer=response.answer,
            sources=sources,
            confidence=response.confidence,
            latency_ms=latency_ms,
            validated=response.validated,
        )

    except Exception as e:
        return QueryResponse(
            success=False,
            error=str(e),
        )


@app.post("/api/search")
async def search_similar(request: SearchRequest):
    """Search for similar chunks."""
    rag = get_rag_system()

    if not rag:
        return {"success": False, "error": "RAG system not available", "results": []}

    try:
        embedder = rag["embedder"]
        store = rag["store"]

        query_embedding = embedder.embed_text(request.query)

        filters = None
        if request.doc_filter:
            filters = {"document_id": request.doc_filter}

        results = store.search(
            query_embedding=query_embedding,
            top_k=request.top_k,
            filters=filters,
        )

        return {
            "success": True,
            "results": [
                {
                    "chunk_id": r.chunk_id,
                    "document_id": r.document_id,
                    "text": r.text,
                    "similarity": r.similarity,
                    "page": r.page,
                    "metadata": r.metadata,
                }
                for r in results
            ]
        }

    except Exception as e:
        return {"success": False, "error": str(e), "results": []}


@app.get("/api/documents", response_model=List[DocumentInfo])
async def list_documents():
    """List all indexed documents."""
    rag = get_rag_system()

    if not rag:
        return []

    try:
        store = rag["store"]
        collection = store._collection

        results = collection.get(include=["metadatas"])
        if not results or not results.get("metadatas"):
            return []

        doc_info = {}
        for meta in results["metadatas"]:
            doc_id = meta.get("document_id", "unknown")
            if doc_id not in doc_info:
                doc_info[doc_id] = {
                    "doc_id": doc_id,
                    "filename": meta.get("source_path", ""),
                    "chunk_count": 0,
                }
            doc_info[doc_id]["chunk_count"] += 1

        return [DocumentInfo(**info) for info in doc_info.values()]

    except Exception as e:
        return []


@app.delete("/api/documents/{doc_id}")
async def delete_document(doc_id: str):
    """Delete a document from the index."""
    rag = get_rag_system()

    if not rag:
        return {"success": False, "error": "RAG system not available"}

    try:
        store = rag["store"]
        collection = store._collection

        # Get chunk IDs for this document
        results = collection.get(
            where={"document_id": doc_id},
            include=[]
        )

        if results and results.get("ids"):
            collection.delete(ids=results["ids"])
            return {"success": True, "deleted_chunks": len(results["ids"])}

        return {"success": False, "error": "Document not found"}

    except Exception as e:
        return {"success": False, "error": str(e)}


# ============================================================================
# Run Server
# ============================================================================

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)