api/main.py

"""
ClauseGuard — FastAPI Backend v4.0
══════════════════════════════════
New in v4.0:
  • /api/redline — clause redlining suggestions
  • /api/chat — RAG chatbot (streaming)
  • /api/ocr — OCR scanned PDF extraction
  • Updated analysis to include redlining data
"""

import os
import re
import json
import time
import uuid
import tempfile
from contextlib import asynccontextmanager
from typing import Optional
from collections import defaultdict
from datetime import datetime

import httpx
import numpy as np
from fastapi import FastAPI, HTTPException, Depends, Body, Request, UploadFile, File as FastAPIFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field

from auth import get_current_user, require_auth

# ── Import shared modules ──
import sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

try:
    from app import (
        split_clauses, classify_cuad, extract_entities,
        detect_contradictions, compute_risk_score, analyze_contract,
        CUAD_LABELS, RISK_MAP, DESC_MAP, _model_status,
        cuad_model, cuad_tokenizer
    )
    from obligations import extract_obligations
    from compliance import check_compliance
    from compare import compare_contracts
    from redlining import generate_redlines
    from chatbot import index_contract, chat_respond
    from ocr_engine import parse_pdf_smart, get_ocr_status
    _SHARED_MODULES = True
except ImportError as e:
    _SHARED_MODULES = False
    print(f"[API] WARNING: Could not import shared modules: {e}")

# ─── Config ───
SUPABASE_URL = os.environ.get("SUPABASE_URL", "")
SUPABASE_SERVICE_KEY = os.environ.get("SUPABASE_SERVICE_ROLE_KEY", "")
HF_API_TOKEN = os.environ.get("HF_API_TOKEN", "")
SAULLM_ENDPOINT = os.environ.get("SAULLM_ENDPOINT", "")
MAX_TEXT_LENGTH = int(os.environ.get("MAX_TEXT_LENGTH", "100000"))

# ─── Rate Limiting ───
_rate_limits = {}
RATE_LIMIT_REQUESTS = 30
RATE_LIMIT_WINDOW = 60

def _check_rate_limit(client_ip: str) -> bool:
    now = time.time()
    if client_ip in _rate_limits:
        count, window_start = _rate_limits[client_ip]
        if now - window_start > RATE_LIMIT_WINDOW:
            _rate_limits[client_ip] = (1, now)
            return True
        if count >= RATE_LIMIT_REQUESTS:
            return False
        _rate_limits[client_ip] = (count + 1, window_start)
        return True
    _rate_limits[client_ip] = (1, now)
    return True

# ─── Supabase helper ───
async def supabase_insert(table: str, data: dict):
    if not SUPABASE_URL or not SUPABASE_SERVICE_KEY:
        return
    try:
        async with httpx.AsyncClient() as client:
            await client.post(
                f"{SUPABASE_URL}/rest/v1/{table}",
                json=data,
                headers={
                    "apikey": SUPABASE_SERVICE_KEY,
                    "Authorization": f"Bearer {SUPABASE_SERVICE_KEY}",
                    "Content-Type": "application/json",
                    "Prefer": "return=minimal",
                },
                timeout=10.0,
            )
    except Exception:
        pass

async def supabase_query(table: str, params: dict, headers_extra: dict = {}):
    if not SUPABASE_URL or not SUPABASE_SERVICE_KEY:
        return []
    try:
        async with httpx.AsyncClient() as client:
            resp = await client.get(
                f"{SUPABASE_URL}/rest/v1/{table}",
                params=params,
                headers={
                    "apikey": SUPABASE_SERVICE_KEY,
                    "Authorization": f"Bearer {SUPABASE_SERVICE_KEY}",
                    **headers_extra,
                },
                timeout=10.0,
            )
            return resp.json() if resp.status_code == 200 else []
    except Exception:
        return []

# ─── In-memory RAG session store ───
_rag_sessions: dict = {}
_RAG_SESSION_MAX = 100

# ─── Request/Response Models ───
class AnalyzeRequest(BaseModel):
    text: Optional[str] = Field(None, min_length=50)
    clauses: Optional[list] = None
    source_url: Optional[str] = None

class CompareRequest(BaseModel):
    text_a: str = Field(..., min_length=50)
    text_b: str = Field(..., min_length=50)

class ExplainRequest(BaseModel):
    clause: str = Field(..., min_length=10, max_length=2000)
    category: str

class ExplainResponse(BaseModel):
    clause: str
    category: str
    explanation: str
    legal_basis: str
    recommendation: str

class ChatRequest(BaseModel):
    message: str = Field(..., min_length=1, max_length=2000)
    session_id: str
    history: Optional[list[dict]] = None

class RedlineRequest(BaseModel):
    session_id: Optional[str] = None
    text: Optional[str] = None
    use_llm: bool = True

# ─── App ───
@asynccontextmanager
async def lifespan(app: FastAPI):
    yield

app = FastAPI(title="ClauseGuard API", version="4.0.0", lifespan=lifespan)

ALLOWED_ORIGINS = [
    "https://clauseguardweb.netlify.app",
    "http://localhost:3000",
    "http://localhost:3001",
]
app.add_middleware(
    CORSMiddleware,
    allow_origins=ALLOWED_ORIGINS,
    allow_origin_regex=r"^chrome-extension://.*$",
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.get("/health")
async def health():
    model_status = "ml" if _SHARED_MODULES and cuad_model else "regex"
    ocr_status = get_ocr_status() if _SHARED_MODULES else "unavailable"
    return {
        "status": "ok",
        "model": model_status,
        "version": "4.0.0",
        "shared_modules": _SHARED_MODULES,
        "ocr": ocr_status,
        "features": ["analyze", "compare", "redline", "chat", "ocr"],
    }

@app.post("/api/analyze")
async def analyze(req: AnalyzeRequest, request: Request, user: Optional[dict] = Depends(get_current_user)):
    client_ip = request.client.host if request.client else "unknown"
    if not _check_rate_limit(client_ip):
        raise HTTPException(status_code=429, detail="Rate limit exceeded.")

    text = req.text
    if not text and req.clauses:
        text = "\n\n".join(req.clauses) if isinstance(req.clauses, list) else str(req.clauses)

    if not text or len(text.strip()) < 50:
        raise HTTPException(status_code=400, detail="Text too short (minimum 50 characters)")
    if len(text) > MAX_TEXT_LENGTH:
        raise HTTPException(status_code=400, detail=f"Text too long (max {MAX_TEXT_LENGTH} chars)")

    start = time.time()

    clauses = split_clauses(text)
    if not clauses:
        raise HTTPException(status_code=400, detail="No clauses detected")

    clause_results = []
    for clause in clauses:
        predictions = classify_cuad(clause)
        if predictions:
            for pred in predictions:
                clause_results.append({
                    "text": clause,
                    "label": pred["label"],
                    "confidence": pred["confidence"],
                    "risk": pred["risk"],
                    "description": pred["description"],
                    "source": pred.get("source", "unknown"),
                })

    entities = extract_entities(text)
    contradictions = detect_contradictions(clause_results, text)
    risk, grade, sev_counts = compute_risk_score(clause_results, len(clauses))
    obligations = extract_obligations(text)
    compliance = check_compliance(text)

    # v4.0: Redlining
    analysis_for_redline = {"clauses": clause_results}
    redlines = []
    try:
        redlines = generate_redlines(analysis_for_redline, use_llm=True)
    except Exception as e:
        print(f"[API] Redlining error: {e}")

    latency = int((time.time() - start) * 1000)

    results_for_db = []
    for cr in clause_results:
        results_for_db.append({
            "text": cr["text"],
            "categories": [{
                "name": cr["label"],
                "severity": cr["risk"],
                "confidence": cr["confidence"],
                "description": cr["description"],
            }],
        })

    # v4.0: RAG indexing
    session_id = None
    try:
        chunks, embeddings, _status = index_contract(text)
        if chunks and embeddings is not None:
            session_id = uuid.uuid4().hex[:12]
            if len(_rag_sessions) >= _RAG_SESSION_MAX:
                oldest = next(iter(_rag_sessions))
                del _rag_sessions[oldest]
            _rag_sessions[session_id] = {
                "chunks": chunks,
                "embeddings": embeddings,
                "analysis": {
                    "risk": {"score": risk, "grade": grade, "breakdown": sev_counts},
                    "metadata": {"total_clauses": len(clauses), "flagged_clauses": len(clause_results)},
                    "clauses": clause_results[:30],
                    "entities": entities[:30],
                    "contradictions": contradictions,
                },
            }
    except Exception as e:
        print(f"[API] RAG indexing error: {e}")

    if user:
        await supabase_insert("analyses", {
            "user_id": user["id"],
            "source_url": req.source_url,
            "total_clauses": len(clauses),
            "flagged_count": len(set(cr["text"] for cr in clause_results)),
            "risk_score": risk,
            "grade": grade,
            "clauses": results_for_db,
            "entities": entities,
            "contradictions": contradictions,
            "obligations": obligations,
            "compliance": compliance,
        })

    return {
        "risk_score": risk,
        "grade": grade,
        "total_clauses": len(clauses),
        "flagged_count": len(set(cr["text"] for cr in clause_results)),
        "results": results_for_db,
        "entities": entities,
        "contradictions": contradictions,
        "obligations": obligations,
        "compliance": compliance,
        "redlines": redlines,
        "model": "ml" if cuad_model else "regex",
        "latency_ms": latency,
        "session_id": session_id,
    }

@app.post("/api/compare")
async def compare(req: CompareRequest, request: Request):
    client_ip = request.client.host if request.client else "unknown"
    if not _check_rate_limit(client_ip):
        raise HTTPException(status_code=429, detail="Rate limit exceeded.")
    return compare_contracts(req.text_a, req.text_b)

@app.post("/api/redline")
async def redline(req: RedlineRequest, request: Request):
    client_ip = request.client.host if request.client else "unknown"
    if not _check_rate_limit(client_ip):
        raise HTTPException(status_code=429, detail="Rate limit exceeded.")

    if req.session_id and req.session_id in _rag_sessions:
        analysis = _rag_sessions[req.session_id]["analysis"]
    elif req.text:
        result, error = analyze_contract(req.text)
        if error:
            raise HTTPException(status_code=400, detail=error)
        analysis = result
    else:
        raise HTTPException(status_code=400, detail="Provide session_id or text")

    redlines = generate_redlines(analysis, use_llm=req.use_llm)
    return {"redlines": redlines, "count": len(redlines)}

@app.post("/api/chat")
async def chat(req: ChatRequest, request: Request):
    client_ip = request.client.host if request.client else "unknown"
    if not _check_rate_limit(client_ip):
        raise HTTPException(status_code=429, detail="Rate limit exceeded.")

    if req.session_id not in _rag_sessions:
        raise HTTPException(status_code=404, detail="Session not found. Analyze a contract first.")

    session = _rag_sessions[req.session_id]
    response_text = ""
    for partial in chat_respond(req.message, req.history or [],
                                 session["chunks"], session["embeddings"], session["analysis"]):
        response_text = partial

    return {"response": response_text, "session_id": req.session_id}

@app.post("/api/chat/stream")
async def chat_stream(req: ChatRequest, request: Request):
    client_ip = request.client.host if request.client else "unknown"
    if not _check_rate_limit(client_ip):
        raise HTTPException(status_code=429, detail="Rate limit exceeded.")

    if req.session_id not in _rag_sessions:
        raise HTTPException(status_code=404, detail="Session not found.")

    session = _rag_sessions[req.session_id]

    async def generate():
        last = ""
        for partial in chat_respond(
            req.message, req.history or [],
            session["chunks"], session["embeddings"], session["analysis"]
        ):
            delta = partial[len(last):]
            last = partial
            if delta:
                yield f"data: {json.dumps({'delta': delta})}\n\n"
        yield "data: [DONE]\n\n"

    return StreamingResponse(generate(), media_type="text/event-stream")

@app.post("/api/ocr")
async def ocr_endpoint(file: UploadFile = FastAPIFile(...)):
    if not file.filename or not file.filename.lower().endswith(".pdf"):
        raise HTTPException(status_code=400, detail="Only PDF files supported")

    with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp:
        content = await file.read()
        tmp.write(content)
        tmp_path = tmp.name

    try:
        text, error, method = parse_pdf_smart(tmp_path)
        if error:
            raise HTTPException(status_code=400, detail=error)
        return {"text": text, "method": method, "chars": len(text) if text else 0, "filename": file.filename}
    finally:
        os.unlink(tmp_path)

@app.post("/api/explain", response_model=ExplainResponse)
async def explain(req: ExplainRequest, user: dict = Depends(require_auth)):
    desc = DESC_MAP.get(req.category, "Unknown category.")
    legal = "Consult local consumer protection laws."
    recommendation = "Review this clause carefully."

    if SAULLM_ENDPOINT and HF_API_TOKEN:
        try:
            prompt = (
                f"Analyze this contract clause and explain why it may be risky.\n\n"
                f"Clause: \"{req.clause}\"\nCategory: {req.category}\n\n"
                f"Provide: 1) Plain-English explanation 2) Legal basis 3) Recommendation"
            )
            async with httpx.AsyncClient(timeout=30.0) as client:
                resp = await client.post(
                    SAULLM_ENDPOINT,
                    json={"inputs": prompt, "parameters": {"max_new_tokens": 300, "temperature": 0.3}},
                    headers={"Authorization": f"Bearer {HF_API_TOKEN}"},
                )
                if resp.status_code == 200:
                    output = resp.json()
                    generated = output[0]["generated_text"] if isinstance(output, list) else output.get("generated_text", "")
                    if generated and len(generated) > 50:
                        parts = generated.split("\n\n")
                        desc = parts[0] if len(parts) > 0 else desc
                        legal = parts[1] if len(parts) > 1 else legal
                        recommendation = parts[2] if len(parts) > 2 else recommendation
        except Exception:
            pass

    return ExplainResponse(clause=req.clause, category=req.category,
                           explanation=desc, legal_basis=legal, recommendation=recommendation)

@app.get("/api/history")
async def history(user: dict = Depends(require_auth), limit: int = 20, offset: int = 0):
    limit = min(limit, 100)
    data = await supabase_query("analyses", {
        "user_id": f"eq.{user['id']}", "select": "*",
        "order": "created_at.desc", "limit": str(limit), "offset": str(offset),
    })
    return {"analyses": data, "limit": limit, "offset": offset}

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