src/finee/api.py

"""
FinEE API - FastAPI Backend
============================

RESTful API for financial entity extraction with:
- Single/batch extraction endpoints
- RAG-enhanced extraction
- PDF/Image processing
- Multi-turn chat
- Analytics

Author: Ranjit Behera
"""

import os
import json
import logging
from datetime import datetime
from typing import List, Dict, Optional, Any
from pathlib import Path

from fastapi import FastAPI, HTTPException, UploadFile, File, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
import uvicorn

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


# ============================================================================
# PYDANTIC MODELS
# ============================================================================

class ExtractionRequest(BaseModel):
    """Single message extraction request."""
    message: str = Field(..., description="Bank SMS or email to extract from")
    use_rag: bool = Field(True, description="Use RAG for context-aware extraction")
    use_llm: bool = Field(False, description="Use LLM for complex cases")


class BatchExtractionRequest(BaseModel):
    """Batch extraction request."""
    messages: List[str] = Field(..., description="List of messages to extract")
    use_rag: bool = True
    use_llm: bool = False


class ExtractionResult(BaseModel):
    """Extraction result."""
    amount: Optional[float] = None
    type: Optional[str] = None
    account: Optional[str] = None
    bank: Optional[str] = None
    date: Optional[str] = None
    time: Optional[str] = None
    reference: Optional[str] = None
    merchant: Optional[str] = None
    beneficiary: Optional[str] = None
    vpa: Optional[str] = None
    category: Optional[str] = None
    is_p2m: Optional[bool] = None
    balance: Optional[float] = None
    status: Optional[str] = None
    confidence: float = 0.0


class ExtractionResponse(BaseModel):
    """API response for extraction."""
    success: bool
    data: Optional[ExtractionResult] = None
    raw_text: Optional[str] = None
    rag_context: Optional[Dict] = None
    processing_time_ms: float = 0
    error: Optional[str] = None


class ChatMessage(BaseModel):
    """Chat message."""
    role: str = Field(..., description="'user' or 'assistant'")
    content: str


class ChatRequest(BaseModel):
    """Chat request for multi-turn analysis."""
    messages: List[ChatMessage]
    context: Optional[Dict] = None


class AnalyticsRequest(BaseModel):
    """Analytics request."""
    transactions: List[Dict]
    period: Optional[str] = "month"
    group_by: Optional[str] = "category"


# ============================================================================
# FASTAPI APP
# ============================================================================

app = FastAPI(
    title="FinEE API",
    description="Financial Entity Extraction API for Indian Banking",
    version="2.0.0",
    docs_url="/docs",
    redoc_url="/redoc",
)

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

# Global state
_extractor = None
_rag_engine = None


def get_extractor():
    """Lazy load extractor."""
    global _extractor
    if _extractor is None:
        try:
            from finee import FinancialExtractor
            _extractor = FinancialExtractor(use_llm=False)
            logger.info("Extractor initialized")
        except ImportError:
            logger.warning("FinEE not installed, using mock extractor")
            _extractor = MockExtractor()
    return _extractor


def get_rag_engine():
    """Lazy load RAG engine."""
    global _rag_engine
    if _rag_engine is None:
        try:
            from finee.rag import RAGEngine
            _rag_engine = RAGEngine()
            logger.info("RAG engine initialized")
        except ImportError:
            logger.warning("RAG not available")
            _rag_engine = None
    return _rag_engine


class MockExtractor:
    """Mock extractor for testing."""
    def extract(self, text: str) -> Dict:
        import re
        result = {}
        
        # Amount
        amount_match = re.search(r'Rs\.?\s*([\d,]+(?:\.\d{2})?)', text)
        if amount_match:
            result['amount'] = float(amount_match.group(1).replace(',', ''))
        
        # Type
        if any(w in text.lower() for w in ['debit', 'debited', 'paid', 'spent']):
            result['type'] = 'debit'
        elif any(w in text.lower() for w in ['credit', 'credited', 'received']):
            result['type'] = 'credit'
        
        return result


# ============================================================================
# ENDPOINTS
# ============================================================================

@app.get("/")
async def root():
    """Health check."""
    return {
        "status": "healthy",
        "service": "FinEE API",
        "version": "2.0.0",
        "timestamp": datetime.utcnow().isoformat()
    }


@app.get("/health")
async def health():
    """Detailed health check."""
    return {
        "status": "healthy",
        "components": {
            "extractor": _extractor is not None,
            "rag": _rag_engine is not None,
        }
    }


@app.post("/extract", response_model=ExtractionResponse)
async def extract(request: ExtractionRequest):
    """
    Extract financial entities from a single message.
    
    - **message**: The bank SMS, email, or notification text
    - **use_rag**: Enable RAG for context-aware extraction
    - **use_llm**: Use LLM for complex cases (slower but more accurate)
    """
    import time
    start = time.time()
    
    try:
        extractor = get_extractor()
        rag = get_rag_engine() if request.use_rag else None
        
        # RAG context
        rag_context = None
        if rag:
            context = rag.retrieve(request.message)
            rag_context = {
                "merchant_info": context.merchant_info,
                "similar_transactions": context.similar_transactions,
                "category_hierarchy": context.category_hierarchy,
            }
        
        # Extract
        result = extractor.extract(request.message)
        
        # Enhance with RAG
        if rag_context and rag_context.get("merchant_info"):
            if not result.get("merchant"):
                result["merchant"] = rag_context["merchant_info"]["name"]
            if not result.get("category"):
                result["category"] = rag_context["merchant_info"]["category"]
            if "is_p2m" not in result:
                result["is_p2m"] = rag_context["merchant_info"]["is_p2m"]
        
        processing_time = (time.time() - start) * 1000
        
        return ExtractionResponse(
            success=True,
            data=ExtractionResult(**result),
            raw_text=request.message,
            rag_context=rag_context,
            processing_time_ms=round(processing_time, 2)
        )
    
    except Exception as e:
        logger.error(f"Extraction failed: {e}")
        return ExtractionResponse(
            success=False,
            error=str(e),
            processing_time_ms=round((time.time() - start) * 1000, 2)
        )


@app.post("/extract/batch")
async def extract_batch(request: BatchExtractionRequest):
    """
    Extract entities from multiple messages.
    
    - **messages**: List of messages to process
    - Returns list of extraction results
    """
    results = []
    
    for message in request.messages:
        req = ExtractionRequest(
            message=message,
            use_rag=request.use_rag,
            use_llm=request.use_llm
        )
        result = await extract(req)
        results.append(result)
    
    return {
        "success": True,
        "total": len(results),
        "successful": sum(1 for r in results if r.success),
        "results": results
    }


@app.post("/parse/pdf")
async def parse_pdf(file: UploadFile = File(...)):
    """
    Parse bank statement PDF and extract transactions.
    
    - **file**: PDF file of bank statement
    - Returns list of extracted transactions
    """
    if not file.filename.endswith('.pdf'):
        raise HTTPException(400, "Only PDF files are supported")
    
    try:
        # Read PDF
        content = await file.read()
        
        # Parse PDF (placeholder - needs pdfplumber)
        transactions = []
        
        # TODO: Implement PDF parsing
        # from pdfplumber import open as open_pdf
        # with open_pdf(io.BytesIO(content)) as pdf:
        #     for page in pdf.pages:
        #         text = page.extract_text()
        #         ...
        
        return {
            "success": True,
            "filename": file.filename,
            "transactions": transactions,
            "message": "PDF parsing not yet implemented"
        }
    
    except Exception as e:
        raise HTTPException(500, f"PDF parsing failed: {e}")


@app.post("/parse/image")
async def parse_image(file: UploadFile = File(...)):
    """
    Parse screenshot/image using OCR and extract entities.
    
    - **file**: Image file (PNG, JPG)
    - Returns extracted text and entities
    """
    allowed = ['.png', '.jpg', '.jpeg', '.webp']
    ext = Path(file.filename).suffix.lower()
    
    if ext not in allowed:
        raise HTTPException(400, f"Only {allowed} files are supported")
    
    try:
        content = await file.read()
        
        # OCR (placeholder - needs pytesseract or EasyOCR)
        extracted_text = ""
        
        # TODO: Implement OCR
        # import pytesseract
        # from PIL import Image
        # image = Image.open(io.BytesIO(content))
        # extracted_text = pytesseract.image_to_string(image)
        
        # Extract entities from OCR text
        if extracted_text:
            extractor = get_extractor()
            result = extractor.extract(extracted_text)
        else:
            result = {}
        
        return {
            "success": True,
            "filename": file.filename,
            "extracted_text": extracted_text,
            "entities": result,
            "message": "Image OCR not yet implemented"
        }
    
    except Exception as e:
        raise HTTPException(500, f"Image parsing failed: {e}")


@app.post("/chat")
async def chat(request: ChatRequest):
    """
    Multi-turn chat for financial analysis.
    
    - **messages**: Conversation history
    - **context**: Optional transaction context
    """
    try:
        # Get last user message
        user_messages = [m for m in request.messages if m.role == "user"]
        if not user_messages:
            raise HTTPException(400, "No user message found")
        
        last_message = user_messages[-1].content
        
        # Simple intent detection
        intent = detect_intent(last_message)
        
        # Generate response based on intent
        response = generate_response(intent, last_message, request.context)
        
        return {
            "success": True,
            "response": response,
            "intent": intent,
        }
    
    except Exception as e:
        raise HTTPException(500, f"Chat failed: {e}")


@app.post("/analytics")
async def analytics(request: AnalyticsRequest):
    """
    Generate spending analytics from transactions.
    
    - **transactions**: List of extracted transactions
    - **period**: Time period (week, month, year)
    - **group_by**: Grouping (category, merchant, type)
    """
    try:
        transactions = request.transactions
        
        if not transactions:
            return {"success": True, "data": {}}
        
        # Group and aggregate
        groups = {}
        total = 0
        
        for txn in transactions:
            key = txn.get(request.group_by, "other")
            amount = txn.get("amount", 0)
            txn_type = txn.get("type", "debit")
            
            if key not in groups:
                groups[key] = {"total": 0, "count": 0, "transactions": []}
            
            if txn_type == "debit":
                groups[key]["total"] += amount
                total += amount
            
            groups[key]["count"] += 1
            groups[key]["transactions"].append(txn)
        
        # Calculate percentages
        for key in groups:
            groups[key]["percentage"] = round(groups[key]["total"] / total * 100, 1) if total > 0 else 0
        
        # Sort by total
        sorted_groups = dict(sorted(groups.items(), key=lambda x: x[1]["total"], reverse=True))
        
        return {
            "success": True,
            "period": request.period,
            "group_by": request.group_by,
            "total_spent": total,
            "transaction_count": len(transactions),
            "breakdown": sorted_groups
        }
    
    except Exception as e:
        raise HTTPException(500, f"Analytics failed: {e}")


@app.get("/merchants")
async def list_merchants(
    category: Optional[str] = None,
    limit: int = 50
):
    """
    List known merchants from knowledge base.
    
    - **category**: Filter by category
    - **limit**: Max results
    """
    rag = get_rag_engine()
    
    if not rag:
        return {"success": False, "error": "RAG not available"}
    
    merchants = []
    for name, merchant in rag.merchant_kb.merchants.items():
        if category and merchant.category != category:
            continue
        
        merchants.append({
            "name": merchant.name,
            "category": merchant.category,
            "vpa": merchant.vpa,
            "is_p2m": merchant.is_p2m,
        })
        
        if len(merchants) >= limit:
            break
    
    return {
        "success": True,
        "count": len(merchants),
        "merchants": merchants
    }


@app.get("/categories")
async def list_categories():
    """List available transaction categories."""
    from finee.rag import CategoryTaxonomy
    
    categories = []
    for name, info in CategoryTaxonomy.TAXONOMY.items():
        categories.append({
            "name": name,
            "parent": info.get("parent"),
            "children": info.get("children", []),
            "keywords": info.get("keywords", [])
        })
    
    return {
        "success": True,
        "categories": categories
    }


# ============================================================================
# HELPER FUNCTIONS
# ============================================================================

def detect_intent(message: str) -> str:
    """Simple intent detection."""
    message_lower = message.lower()
    
    if any(w in message_lower for w in ['how much', 'total', 'spent', 'spending']):
        return "spending_query"
    elif any(w in message_lower for w in ['compare', 'vs', 'versus', 'difference']):
        return "comparison"
    elif any(w in message_lower for w in ['category', 'break', 'breakdown']):
        return "category_breakdown"
    elif any(w in message_lower for w in ['extract', 'parse', 'analyze']):
        return "extraction"
    else:
        return "general"


def generate_response(intent: str, message: str, context: Optional[Dict]) -> str:
    """Generate chat response based on intent."""
    if intent == "spending_query":
        if context and "transactions" in context:
            total = sum(t.get("amount", 0) for t in context["transactions"] if t.get("type") == "debit")
            return f"Based on your transactions, you've spent ₹{total:,.2f}"
        return "Please share your transaction data for spending analysis."
    
    elif intent == "category_breakdown":
        return "I can break down your spending by category. Please share transaction data."
    
    elif intent == "comparison":
        return "To compare periods, please specify the time ranges you'd like to compare."
    
    elif intent == "extraction":
        return "Share a bank message and I'll extract the financial details."
    
    else:
        return "I can help you analyze transactions, extract entities, or provide spending insights. What would you like to know?"


# ============================================================================
# MAIN
# ============================================================================

def start_server(host: str = "0.0.0.0", port: int = 8000):
    """Start the API server."""
    uvicorn.run(app, host=host, port=port)


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="FinEE API Server")
    parser.add_argument("--host", default="0.0.0.0", help="Host to bind")
    parser.add_argument("--port", type=int, default=8000, help="Port to bind")
    parser.add_argument("--reload", action="store_true", help="Enable auto-reload")
    
    args = parser.parse_args()
    
    if args.reload:
        uvicorn.run("api:app", host=args.host, port=args.port, reload=True)
    else:
        start_server(args.host, args.port)