api_server.py

"""
FastAPI Backend Server
========================
REST API for the Indian Cyber Law RAG Chatbot.
Supports multilingual queries and PDF ingestion.
"""

import os
import shutil
import tempfile
from pathlib import Path
from typing import Optional
from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException, UploadFile, File, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
from loguru import logger
from dotenv import load_dotenv

load_dotenv()

# ─────────────────────────────────────────────────────────────────
# Startup / Shutdown
# ─────────────────────────────────────────────────────────────────

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Initialize the RAG pipeline on startup."""
    logger.info("🚀 Starting Indian Cyber Law RAG Server...")
    global rag_pipeline
    from utils.rag_pipeline import get_rag_pipeline
    rag_pipeline = get_rag_pipeline()
    logger.success("✅ Server ready!")
    yield
    logger.info("Server shutting down...")


# ─────────────────────────────────────────────────────────────────
# App Setup
# ─────────────────────────────────────────────────────────────────

app = FastAPI(
    title="Indian Cyber Law Legal Advisor API",
    description="""
    AI-powered legal advisor for Indian cyber laws including:
    - IT Act 2000 & IT Amendment Act 2008
    - Digital Personal Data Protection Act 2023
    - IT Rules 2021
    - CERT-In Guidelines
    - Related IPC sections and case laws
    
    Supports 12 Indian languages!
    """,
    version="1.0.0",
    lifespan=lifespan,
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Restrict in production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

rag_pipeline = None


# ─────────────────────────────────────────────────────────────────
# Request / Response Models
# ─────────────────────────────────────────────────────────────────

class QueryRequest(BaseModel):
    query: str = Field(..., min_length=3, max_length=2000,
                       description="Legal query in any supported Indian language")
    language: Optional[str] = Field(None, description="Force language (e.g., 'hi', 'ta'). Auto-detected if None.")
    top_k: int = Field(6, ge=1, le=15, description="Number of documents to retrieve")
    category_filter: Optional[str] = Field(None,
        description="Filter: criminal_offence|data_protection|privacy|cyber_fraud|identity_theft|intermediary_liability|constitutional_law")
    translate_response: bool = Field(True, description="Translate response to user's language")
    session_id: Optional[str] = Field(None, description="Session ID for conversation continuity")
    mode: str = Field("general", description="Prompt mode (e.g., general, complaint, document, article)")

class SourceCitation(BaseModel):
    source: str
    section: str
    type: str
    category: str
    score: float
    snippet: str

class QueryResponse(BaseModel):
    answer: str
    sources: list[SourceCitation]
    detected_language: str
    language_name: str
    translated: bool
    confidence: str
    query_english: str
    chunks_retrieved: int
    session_id: Optional[str]

class HealthResponse(BaseModel):
    status: str
    documents_in_store: int
    supported_languages: list[str]
    llm_provider: str
    model: str

class IngestResponse(BaseModel):
    status: str
    documents_added: int
    message: str


# ─────────────────────────────────────────────────────────────────
# Session Store (in-memory; use Redis in production)
# ─────────────────────────────────────────────────────────────────

session_histories: dict[str, list] = {}


# ─────────────────────────────────────────────────────────────────
# Endpoints
# ─────────────────────────────────────────────────────────────────

@app.get("/", tags=["Info"])
async def root():
    return {
        "name": "Indian Cyber Law Legal Advisor",
        "version": "1.0.0",
        "docs": "/docs",
        "health": "/health",
        "endpoints": {
            "query":     "POST /query",
            "ingest_pdf": "POST /ingest/pdf",
            "languages": "GET /languages",
            "stats":     "GET /stats",
        }
    }


@app.get("/health", response_model=HealthResponse, tags=["Info"])
async def health_check():
    """Check API health and vector store status."""
    from utils.multilingual import SUPPORTED_LANGUAGES

    store_count = rag_pipeline.vector_store.collection.count() if rag_pipeline else 0

    return HealthResponse(
        status="healthy" if rag_pipeline else "initializing",
        documents_in_store=store_count,
        supported_languages=list(SUPPORTED_LANGUAGES.keys()),
        llm_provider=os.getenv("LLM_PROVIDER", "openai"),
        model=os.getenv("LLM_MODEL", "gpt-4o"),
    )


@app.get("/languages", tags=["Info"])
async def get_supported_languages():
    """Get list of supported languages for the chatbot."""
    from utils.multilingual import SUPPORTED_LANGUAGES
    return {
        "supported_languages": [
            {
                "code":   code,
                "name":   info["name"],
                "native": info["native"],
                "script": info["script"],
            }
            for code, info in SUPPORTED_LANGUAGES.items()
        ]
    }


@app.get("/stats", tags=["Info"])
async def get_stats():
    """Get vector store statistics."""
    if not rag_pipeline:
        raise HTTPException(status_code=503, detail="RAG pipeline not initialized")

    stats = rag_pipeline.vector_store.get_collection_stats()
    return {"vector_store": stats}


@app.post("/query", response_model=QueryResponse, tags=["Query"])
async def query_cyber_law(request: QueryRequest):
    """
    Query the Indian Cyber Law knowledge base.
    
    Supports queries in English and 11 Indian languages.
    Returns legal information with citations.
    
    **Example queries:**
    - "What is the punishment for hacking under IT Act?"
    - "साइबर क्राइम की शिकायत कैसे करें?"
    - "What are my rights if my data is breached?"
    - "Section 66C IT Act क्या है?"
    """
    if not rag_pipeline:
        raise HTTPException(status_code=503, detail="RAG pipeline not initialized")

    try:
        # Apply session history if session_id provided
        if request.session_id and request.session_id in session_histories:
            rag_pipeline.conversation_history = session_histories[request.session_id]

        result = rag_pipeline.query(
            user_query=request.query,
            top_k=request.top_k,
            category_filter=request.category_filter,
            user_language=request.language,
            translate_response=request.translate_response,
            mode=request.mode,
        )

        # Save session history
        if request.session_id:
            session_histories[request.session_id] = rag_pipeline.conversation_history.copy()

        # Get language name
        from utils.multilingual import SUPPORTED_LANGUAGES
        lang_code = result["detected_language"]
        lang_name = SUPPORTED_LANGUAGES.get(lang_code, {}).get("name", "Unknown")

        return QueryResponse(
            answer=result["answer"],
            sources=[SourceCitation(**s) for s in result["sources"]],
            detected_language=lang_code,
            language_name=lang_name,
            translated=result["translated"],
            confidence=result["confidence"],
            query_english=result["query_english"],
            chunks_retrieved=result["chunks_retrieved"],
            session_id=request.session_id,
        )

    except Exception as e:
        logger.error("Query failed: {}", e)
        raise HTTPException(status_code=500, detail=f"Query processing failed: {str(e)}")


@app.post("/query/clear-history", tags=["Query"])
async def clear_session_history(session_id: Optional[str] = None):
    """Clear conversation history for a session."""
    if session_id and session_id in session_histories:
        del session_histories[session_id]

    rag_pipeline.clear_history()
    return {"status": "success", "message": "Conversation history cleared"}


@app.post("/ingest/pdf", response_model=IngestResponse, tags=["Ingestion"])
async def ingest_pdf(
    file: UploadFile = File(..., description="PDF file (IT Act, court judgment, RBI circular, etc.)"),
    source: str = "Custom Document",
    category: str = "legal_document",
    background_tasks: BackgroundTasks = None,
):
    """
    Upload and ingest a PDF document into the knowledge base.
    
    Accepts: IT Act PDFs, court judgment PDFs, RBI/SEBI circulars, CERT-In advisories, etc.
    """
    if not file.filename.endswith(".pdf"):
        raise HTTPException(status_code=400, detail="Only PDF files are accepted")

    if not rag_pipeline:
        raise HTTPException(status_code=503, detail="RAG pipeline not initialized")

    # Save uploaded file temporarily
    with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp:
        shutil.copyfileobj(file.file, tmp)
        tmp_path = tmp.name

    try:
        metadata = {
            "source":   source,
            "section":  "auto-chunked",
            "type":     "pdf_document",
            "category": category,
            "jurisdiction": "India",
        }
        added = rag_pipeline.add_pdf_document(tmp_path, metadata)
        return IngestResponse(
            status="success",
            documents_added=added,
            message=f"Successfully ingested {added} document chunks from '{file.filename}'"
        )
    except Exception as e:
        logger.error("PDF ingestion failed: {}", e)
        raise HTTPException(status_code=500, detail=f"PDF ingestion failed: {str(e)}")
    finally:
        os.unlink(tmp_path)


# ─────────────────────────────────────────────────────────────────
# Entry Point
# ─────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(
        "api_server:app",
        host=os.getenv("APP_HOST", "0.0.0.0"),
        port=int(os.getenv("APP_PORT", "8000")),
        reload=os.getenv("DEBUG", "false").lower() == "true",
        log_level=os.getenv("LOG_LEVEL", "info").lower(),
    )