"""
main.py — Rwanda Legal AI · FastAPI Backend (v8)
=================================================

Production-oriented FastAPI backend for Amategeko AI.

Features
--------
POST   /api/chat
    Full legal AI response:
    - answer
    - sources
    - confidence
    - citation fidelity
    - situation / agent decision
    - actions
    - lawyer recommendations
    - uploaded document context

POST   /api/upload
    Parses:
    - PDF
    - DOCX
    - DOC
    - CSV
    - XLSX
    - XLS
    - TXT
    - MD

GET    /api/health
    System + vector database status.

GET    /api/law/{law_id}
    Retrieves indexed law information and articles.

DELETE /api/session/{session_id}
    Clears conversation memory.

DELETE /api/session/{session_id}/file
    Clears frontend/backend file context if session storage is used.

GET    /api/sessions
    Development/debug endpoint.

POST   /api/ingest
    Manually triggers Gazette ingestion.

Architecture
------------
Frontend
   ↓
FastAPI
   ↓
RwandaLegalAI
   ├── Query analysis
   ├── Hybrid legal retrieval
   ├── RAG
   ├── Citation validation
   ├── Situation classification
   ├── Action Agent
   └── Lawyer matching
           ↓
Official Gazette / ChromaDB
"""

import os
import re
import time
import uuid
import logging
import warnings

from typing import Any, Dict, List, Optional

from fastapi import FastAPI, HTTPException, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware

from pydantic import BaseModel, Field, ConfigDict

from dotenv import load_dotenv
import uvicorn


# ─────────────────────────────────────────────
# PROJECT IMPORTS
# ─────────────────────────────────────────────

from services.vector_store import get_vectorstore

from app.chat import RwandaLegalAI

from services.file_parser import parse_uploaded_file


# ─────────────────────────────────────────────
# OPTIONAL / CANONICAL INGESTION
# ─────────────────────────────────────────────

try:
    from force_ingest import ingest_all_documents
except ImportError:
    ingest_all_documents = None


# ─────────────────────────────────────────────
# SETUP
# ─────────────────────────────────────────────

load_dotenv()

logging.basicConfig(
    level=logging.INFO,
    format="%(levelname)s │ %(message)s",
)

logger = logging.getLogger(__name__)

warnings.filterwarnings("ignore")

os.environ["TRANSFORMERS_VERBOSITY"] = "error"
os.environ["TOKENIZERS_PARALLELISM"] = "false"


# ─────────────────────────────────────────────
# APP
# ─────────────────────────────────────────────

app = FastAPI(
    title="Rwanda Legal AI API",
    description=(
        "Official Gazette-backed legal intelligence platform "
        "for Rwanda. Supports legal research, document analysis, "
        "action guidance, and lawyer matching."
    ),
    version="8.0.0",
)


# ─────────────────────────────────────────────
# CORS
# ─────────────────────────────────────────────

FRONTEND_URL = os.getenv(
    "FRONTEND_URL",
    "https://your-production-domain.com",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "http://localhost:3000",
        "http://localhost:3001",
        FRONTEND_URL,
    ],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# ═════════════════════════════════════════════
# FILE UPLOAD CONFIGURATION
# ═════════════════════════════════════════════

ALLOWED_EXTENSIONS = {
    "pdf",
    "docx",
    "doc",
    "csv",
    "xlsx",
    "xls",
    "txt",
    "md",
}

MAX_UPLOAD_BYTES = 20 * 1024 * 1024

MAX_FILE_CONTEXT_CHARS = 24000


# ═════════════════════════════════════════════
# PYDANTIC MODELS
# ═════════════════════════════════════════════

class ChatRequest(BaseModel):
    """
    Incoming chat request.

    file_context and file_summary are optional so normal
    legal questions continue to work without uploads.
    """

    question: str = Field(
        ...,
        min_length=1,
        max_length=2000,
    )

    session_id: str = Field(
        default="default",
        min_length=1,
        max_length=200,
    )

    file_context: Optional[str] = Field(
        default=None,
        max_length=MAX_FILE_CONTEXT_CHARS,
    )

    file_summary: Optional[str] = Field(
        default=None,
        max_length=500,
    )


class UploadResponse(BaseModel):
    file_id: str
    filename: str
    summary: str
    file_type: str
    truncated: bool
    text: str


class SourceReference(BaseModel):
    title: str = ""
    article: str = ""
    year: str = ""
    source: str = ""
    excerpt: str = ""
    category: str = ""
    law_id: str = ""


class RecommendedLawyer(BaseModel):
    """
    Rich lawyer recommendation model.

    Supports both the older lawyer schema and the newer
    simplified matching schema.
    """

    id: Optional[str] = None

    name: Optional[str] = None

    full_name: Optional[str] = None

    rba_number: Optional[str] = None

    district: Optional[str] = None

    # New + old naming compatibility
    specialty: Optional[str] = None

    specialties: List[str] = Field(
        default_factory=list
    )

    languages: List[str] = Field(
        default_factory=list
    )

    years_experience: Optional[int] = None

    years_of_experience: Optional[int] = None

    experience_years: Optional[int] = None

    phone: Optional[str] = None

    whatsapp: Optional[str] = None

    rating: Optional[float] = None

    total_cases: Optional[int] = None

    match_score: Optional[float] = None

    match_reasons: List[str] = Field(
        default_factory=list
    )

    bio_summary: Optional[str] = None

    bio: Optional[str] = None

    accepting_clients: Optional[bool] = True

    consultation_fee_rwf: Optional[int] = None

    availability_note: Optional[str] = None

    intro_message: Optional[str] = None

    model_config = ConfigDict(
        populate_by_name=True
    )


class AgentDecision(BaseModel):
    situation: str = "SELF_RESOLVABLE"

    urgency: str = "NORMAL"

    rationale: Optional[str] = None

    citizen_guidance: Optional[str] = None

    how_to_work_with_lawyer: List[str] = Field(
        default_factory=list
    )


class ActionsPayload(BaseModel):
    """
    Action Agent output.
    """

    next_steps: List[Dict[str, str]] = Field(
        default_factory=list
    )

    docs_required: List[str] = Field(
        default_factory=list
    )

    dispute_summary: str = ""

    briefing_note: str = ""

    lawyer_match: Optional[Dict[str, Any]] = None


class ChatResponse(BaseModel):
    """
    Unified response contract combining the old and v7 APIs.
    """

    response: str

    session_id: str

    sources: List[SourceReference] = Field(
        default_factory=list
    )

    confidence: float = 0.0

    confidence_band: str = "LOW"

    # Restored from previous API
    citation_fidelity: float = 0.0

    language: str = "english"

    category: str = "general"

    situation: str = "SELF_RESOLVABLE"

    # Agent decision
    decision: Optional[AgentDecision] = None

    # Action Agent
    actions: ActionsPayload = Field(
        default_factory=ActionsPayload
    )

    # Convenience frontend fields
    howToWorkWithLawyer: List[str] = Field(
        default_factory=list
    )

    recommendedLawyers: List[RecommendedLawyer] = Field(
        default_factory=list
    )

    is_clarifying: bool = False


class HealthResponse(BaseModel):
    status: str

    database_ready: bool

    assistant_ready: bool

    active_sessions: int

    chunk_count: int

    uptime_seconds: float


# ═════════════════════════════════════════════
# GLOBAL STATE
# ═════════════════════════════════════════════

"""
One RwandaLegalAI instance is shared by the API.

Session isolation MUST be handled inside RwandaLegalAI
using session_id.

This is preferable to creating a heavy RAG/embedding
pipeline for every HTTP session.
"""

agent: Optional[RwandaLegalAI] = None

assistant_ready: bool = False

start_time: float = time.time()


# Session metadata.

# IMPORTANT:
# Conversation memory itself should live inside RwandaLegalAI.
# This dictionary only tracks API-level session activity.

sessions: Dict[str, Dict[str, Any]] = {}

SESSION_TTL = 3600


# Uploaded file metadata.

# The actual extracted text is intentionally NOT persisted
# here because the frontend sends file_context with chat.
#
# This dictionary is mainly useful for lifecycle tracking
# and future server-side file storage.

uploaded_files: Dict[str, Dict[str, Any]] = {}


# ═════════════════════════════════════════════
# SESSION MANAGEMENT
# ═════════════════════════════════════════════

def cleanup_stale_sessions() -> None:
    """
    Remove API session metadata that has been inactive
    for SESSION_TTL seconds.
    """

    now = time.time()

    stale_ids = [
        session_id
        for session_id, data in sessions.items()
        if now - data["last_used"] > SESSION_TTL
    ]

    for session_id in stale_ids:
        logger.info(
            "Evicting stale API session: %s",
            session_id,
        )

        sessions.pop(session_id, None)


def touch_session(session_id: str) -> None:
    """
    Create/update API-level session metadata.
    """

    now = time.time()

    cleanup_stale_sessions()

    if session_id not in sessions:
        sessions[session_id] = {
            "created_at": now,
            "last_used": now,
        }
    else:
        sessions[session_id]["last_used"] = now


# ═════════════════════════════════════════════
# STARTUP
# ═════════════════════════════════════════════

@app.on_event("startup")
async def startup_event():
    """
    Initialize the vector database and legal AI agent.
    """

    global agent
    global assistant_ready

    logger.info("Starting Rwanda Legal AI v8...")

    # ─────────────────────────────────────────
    # VECTOR DATABASE
    # ─────────────────────────────────────────

    if not os.path.exists("./chroma_db"):

        logger.warning(
            "No vector DB found — starting legal document ingestion..."
        )

        try:

            if ingest_all_documents is not None:

                # Canonical legal ingestion pipeline.
                ingest_all_documents()

            else:

                # Compatibility fallback.
                from services.vector_store import ingest_gazettes

                ingest_gazettes()

            logger.info(
                "✅ Legal document ingestion complete"
            )

        except Exception as exc:

            logger.error(
                "❌ Ingestion failed: %s",
                exc,
                exc_info=True,
            )

            raise RuntimeError(
                "Database initialization failed. "
                "Check data/gazette/ and ingestion configuration."
            ) from exc

    else:

        try:

            vectorstore = get_vectorstore()

            count = vectorstore._collection.count()

            logger.info(
                "✅ Vector DB found — %,d chunks ready",
                count,
            )

        except Exception as exc:

            logger.warning(
                "Vector DB found but chunk count unavailable: %s",
                exc,
            )

    # ─────────────────────────────────────────
    # AI AGENT
    # ─────────────────────────────────────────

    try:

        agent = RwandaLegalAI()

        assistant_ready = True

        logger.info(
            "🚀 Rwanda Legal AI v8 ready on :8000"
        )

    except Exception as exc:

        logger.error(
            "❌ Failed to initialize RwandaLegalAI: %s",
            exc,
            exc_info=True,
        )

        assistant_ready = False

        raise RuntimeError(
            "RwandaLegalAI initialization failed."
        ) from exc


# ═════════════════════════════════════════════
# HEALTH
# ═════════════════════════════════════════════

@app.get(
    "/api/health",
    response_model=HealthResponse,
)
async def health_check():

    chunk_count = 0

    database_ready = os.path.exists("./chroma_db")

    if database_ready:

        try:

            chunk_count = (
                get_vectorstore()
                ._collection
                .count()
            )

        except Exception as exc:

            logger.warning(
                "Could not read ChromaDB count: %s",
                exc,
            )

    return HealthResponse(

        status=(
            "healthy"
            if assistant_ready
            else "initializing"
        ),

        database_ready=database_ready,

        assistant_ready=assistant_ready,

        active_sessions=len(sessions),

        chunk_count=chunk_count,

        uptime_seconds=round(
            time.time() - start_time,
            1,
        ),
    )


# ═════════════════════════════════════════════
# FILE UPLOAD
# ═════════════════════════════════════════════

@app.post(
    "/api/upload",
    response_model=UploadResponse,
)
async def upload_file(
    file: UploadFile = File(...)
):
    """
    Parse an uploaded legal/personal document.

    The extracted text is returned to the frontend.
    The frontend can then send that text as file_context
    with subsequent /api/chat requests.
    """

    filename = file.filename or "upload"

    # ─────────────────────────────────────────
    # EXTENSION
    # ─────────────────────────────────────────

    ext = (
        filename.rsplit(".", 1)[-1].lower()
        if "." in filename
        else ""
    )

    if ext not in ALLOWED_EXTENSIONS:

        raise HTTPException(
            status_code=415,
            detail=(
                f"File type '.{ext}' is not supported. "
                f"Allowed: "
                f"{', '.join(sorted(ALLOWED_EXTENSIONS))}"
            ),
        )

    # ─────────────────────────────────────────
    # READ
    # ─────────────────────────────────────────

    try:

        file_bytes = await file.read()

    except Exception as exc:

        logger.error(
            "Failed reading uploaded file: %s",
            exc,
            exc_info=True,
        )

        raise HTTPException(
            status_code=400,
            detail="Could not read uploaded file.",
        )

    # ─────────────────────────────────────────
    # SIZE
    # ─────────────────────────────────────────

    if len(file_bytes) > MAX_UPLOAD_BYTES:

        raise HTTPException(
            status_code=413,
            detail=(
                f"File too large "
                f"({len(file_bytes) / 1024 / 1024:.1f} MB). "
                f"Maximum is 20 MB."
            ),
        )

    if len(file_bytes) == 0:

        raise HTTPException(
            status_code=400,
            detail="Uploaded file is empty.",
        )

    # ─────────────────────────────────────────
    # PARSE
    # ─────────────────────────────────────────

    logger.info(
        "Parsing upload: %s (%d bytes)",
        filename,
        len(file_bytes),
    )

    try:

        parsed = parse_uploaded_file(
            file_bytes,
            filename,
        )

    except Exception as exc:

        logger.error(
            "File parser error: %s",
            exc,
            exc_info=True,
        )

        raise HTTPException(
            status_code=422,
            detail=(
                "The uploaded file could not be parsed."
            ),
        )

    if parsed.error:

        raise HTTPException(
            status_code=422,
            detail=parsed.error,
        )

    if not parsed.text or not parsed.text.strip():

        raise HTTPException(
            status_code=422,
            detail=(
                f"No readable text could be extracted "
                f"from '{filename}'. "
                "The file may be scanned without OCR, "
                "password-protected, or empty."
            ),
        )

    # ─────────────────────────────────────────
    # FILE ID
    # ─────────────────────────────────────────

    file_id = str(uuid.uuid4())

    uploaded_files[file_id] = {
        "filename": filename,
        "file_type": parsed.file_type,
        "summary": parsed.summary,
        "created_at": time.time(),
    }

    logger.info(
        "Parsed %s → %,d chars | %s",
        filename,
        len(parsed.text),
        parsed.summary,
    )

    return UploadResponse(

        file_id=file_id,

        filename=filename,

        summary=parsed.summary,

        file_type=parsed.file_type,

        truncated=parsed.truncated,

        text=parsed.text,
    )


# ═════════════════════════════════════════════
# CHAT
# ═════════════════════════════════════════════

@app.post(
    "/api/chat",
    response_model=ChatResponse,
)
async def chat_endpoint(
    request: ChatRequest
):

    if not assistant_ready or agent is None:

        raise HTTPException(
            status_code=503,
            detail=(
                "Assistant is still initializing. "
                "Please retry in a moment."
            ),
        )

    try:

        # ─────────────────────────────────────
        # SESSION
        # ─────────────────────────────────────

        touch_session(
            request.session_id
        )

        # ─────────────────────────────────────
        # SANITIZE FILE CONTEXT
        # ─────────────────────────────────────

        file_context = request.file_context

        if file_context:

            file_context = file_context.strip()

            if len(file_context) > MAX_FILE_CONTEXT_CHARS:

                file_context = (
                    file_context[
                        :MAX_FILE_CONTEXT_CHARS
                    ]
                )

                logger.info(
                    "File context truncated for session %s",
                    request.session_id,
                )

        # ─────────────────────────────────────
        # ASK AGENT
        # ─────────────────────────────────────

        result: Dict[str, Any] = agent.ask_full(

            question=request.question,

            session_id=request.session_id,

            file_context=file_context,

            file_summary=request.file_summary,
        )

        # ─────────────────────────────────────
        # SOURCES
        # ─────────────────────────────────────

        sources: List[SourceReference] = []

        for source in result.get("sources", []):

            if not isinstance(source, dict):
                continue

            sources.append(
                SourceReference(
                    title=str(
                        source.get(
                            "title",
                            "",
                        )
                    ),

                    article=str(
                        source.get(
                            "article",
                            "",
                        )
                    ),

                    year=str(
                        source.get(
                            "year",
                            "",
                        )
                    ),

                    source=str(
                        source.get(
                            "source",
                            "",
                        )
                    ),

                    excerpt=str(
                        source.get(
                            "excerpt",
                            "",
                        )
                    ),

                    category=str(
                        source.get(
                            "category",
                            "",
                        )
                    ),

                    law_id=str(
                        source.get(
                            "law_id",
                            "",
                        )
                    ),
                )
            )

        # ─────────────────────────────────────
        # DECISION
        # ─────────────────────────────────────

        raw_decision = (
            result.get("decision")
            or {}
        )

        if not isinstance(
            raw_decision,
            dict,
        ):
            raw_decision = {}

        decision = (
            AgentDecision(
                situation=raw_decision.get(
                    "situation",
                    result.get(
                        "situation",
                        "SELF_RESOLVABLE",
                    ),
                ),

                urgency=raw_decision.get(
                    "urgency",
                    "NORMAL",
                ),

                rationale=raw_decision.get(
                    "rationale"
                ),

                citizen_guidance=raw_decision.get(
                    "citizen_guidance"
                ),

                how_to_work_with_lawyer=(
                    raw_decision.get(
                        "how_to_work_with_lawyer",
                        [],
                    )
                    or []
                ),
            )
        )

        # ─────────────────────────────────────
        # ACTIONS
        # ─────────────────────────────────────

        raw_actions = (
            result.get("actions")
            or {}
        )

        if not isinstance(
            raw_actions,
            dict,
        ):
            raw_actions = {}

        raw_lawyer_match = (
            raw_actions.get(
                "lawyer_match"
            )
            or {}
        )

        if not isinstance(
            raw_lawyer_match,
            dict,
        ):
            raw_lawyer_match = {}

        how_to_work = (

            raw_lawyer_match.get(
                "how_to_work_with_lawyer"
            )

            or raw_decision.get(
                "how_to_work_with_lawyer"
            )

            or result.get(
                "howToWorkWithLawyer"
            )

            or []

        )

        actions = ActionsPayload(

            next_steps=(
                raw_actions.get(
                    "next_steps",
                    [],
                )
                or []
            ),

            docs_required=(
                raw_actions.get(
                    "docs_required",
                    [],
                )
                or []
            ),

            dispute_summary=str(
                raw_actions.get(
                    "dispute_summary",
                    "",
                )
                or ""
            ),

            briefing_note=str(
                raw_actions.get(
                    "briefing_note",
                    "",
                )
                or ""
            ),

            lawyer_match=(
                raw_lawyer_match
                if raw_lawyer_match
                else None
            ),
        )

        # ─────────────────────────────────────
        # LAWYER MATCHES
        # ─────────────────────────────────────

        recommended: List[
            RecommendedLawyer
        ] = []

        matches = (
            raw_lawyer_match.get(
                "matches",
                [],
            )
            or []
        )

        for lawyer in matches:

            if not isinstance(
                lawyer,
                dict,
            ):
                continue

            years = lawyer.get(
                "experience_years",
                lawyer.get(
                    "years_experience",
                    lawyer.get("years_of_experience"),
                ),
            )

            specialties = (
                lawyer.get(
                    "specialties"
                )
                or []
            )

            if isinstance(
                specialties,
                str,
            ):
                specialties = [
                    specialties
                ]

            languages = (
                lawyer.get(
                    "languages"
                )
                or []
            )

            if isinstance(
                languages,
                str,
            ):
                languages = [
                    languages
                ]

            match_reasons = (
                lawyer.get(
                    "match_reasons"
                )
                or []
            )

            if isinstance(
                match_reasons,
                str,
            ):
                match_reasons = [
                    match_reasons
                ]

            recommended.append(
                RecommendedLawyer(

                    id=(
                        str(
                            lawyer["id"]
                        )
                        if lawyer.get("id")
                        is not None
                        else None
                    ),

                    name=(
                        lawyer.get("name")
                        or lawyer.get(
                            "full_name"
                        )
                    ),

                    full_name=(
                        lawyer.get(
                            "full_name"
                        )
                        or lawyer.get("name")
                    ),

                    rba_number=lawyer.get(
                        "rba_number"
                    ),

                    district=lawyer.get(
                        "district"
                    ),

                    specialty=lawyer.get(
                        "specialty"
                    ),

                    specialties=specialties,

                    languages=languages,

                    years_experience=(
                        int(years)
                        if years is not None
                        else None
                    ),

                    experience_years=(
                        int(years)
                        if years is not None
                        else None
                    ),

                    years_of_experience=(
                        int(years)
                        if years is not None
                        else None
                    ),

                    phone=lawyer.get(
                        "phone"
                    ),

                    whatsapp=lawyer.get(
                        "whatsapp"
                    ),

                    rating=(
                        float(
                            lawyer["rating"]
                        )
                        if lawyer.get(
                            "rating"
                        ) is not None
                        else None
                    ),

                    total_cases=(
                        int(
                            lawyer["total_cases"]
                        )
                        if lawyer.get(
                            "total_cases"
                        ) is not None
                        else None
                    ),

                    match_score=(
                        float(
                            lawyer["match_score"]
                        )
                        if lawyer.get(
                            "match_score"
                        ) is not None
                        else None
                    ),

                    match_reasons=match_reasons,

                    bio_summary=(
                        lawyer.get(
                            "bio_summary"
                        )
                        or lawyer.get(
                            "bio"
                        )
                    ),

                    bio=(
                        lawyer.get("bio")
                        or lawyer.get(
                            "bio_summary"
                        )
                    ),

                    accepting_clients=lawyer.get(
                        "accepting_leads",
                        lawyer.get("accepting_clients", True),
                    ),

                    consultation_fee_rwf=(
                        int(
                            lawyer[
                                "consultation_fee_rwf"
                            ]
                        )
                        if lawyer.get(
                            "consultation_fee_rwf"
                        ) is not None
                        else None
                    ),

                    availability_note=lawyer.get(
                        "availability_note"
                    ),

                    intro_message=(
                        lawyer.get(
                            "intro_message"
                        )
                        or raw_lawyer_match.get(
                            "intro_message"
                        )
                    ),
                )
            )

        # ─────────────────────────────────────
        # RESPONSE
        # ─────────────────────────────────────

        situation = result.get(
            "situation",
            raw_decision.get(
                "situation",
                "SELF_RESOLVABLE",
            ),
        )

        return ChatResponse(

            response=str(
                result.get(
                    "response",
                    "",
                )
            ),

            session_id=request.session_id,

            sources=sources,

            confidence=float(
                result.get(
                    "confidence",
                    0.0,
                )
                or 0.0
            ),

            confidence_band=str(
                result.get(
                    "confidence_band",
                    "LOW",
                )
                or "LOW"
            ),

            citation_fidelity=float(
                result.get(
                    "citation_fidelity",
                    0.0,
                )
                or 0.0
            ),

            language=str(
                result.get(
                    "language",
                    "english",
                )
                or "english"
            ),

            category=str(
                result.get(
                    "category",
                    "general",
                )
                or "general"
            ),

            situation=situation,

            decision=decision,

            actions=actions,

            howToWorkWithLawyer=how_to_work,

            recommendedLawyers=recommended,

            is_clarifying=bool(
                result.get(
                    "is_clarifying",
                    False,
                )
            ),
        )

    except HTTPException:

        raise

    except TypeError as exc:

        # Helpful diagnostic for a mismatched chat.py
        logger.error(
            "RwandaLegalAI.ask_full signature mismatch: %s",
            exc,
            exc_info=True,
        )

        raise HTTPException(
            status_code=500,
            detail=(
                "The AI backend and chat engine are using "
                "different API versions. Check RwandaLegalAI.ask_full()."
            ),
        )

    except Exception as exc:

        logger.error(
            "Chat endpoint error: %s",
            exc,
            exc_info=True,
        )

        raise HTTPException(
            status_code=500,
            detail=(
                "The AI assistant encountered an error. "
                "Please try again."
            ),
        )


# ═════════════════════════════════════════════
# LAW DETAIL
# ═════════════════════════════════════════════

@app.get("/api/law/{law_id}")
async def get_law_detail(
    law_id: str
):
    """
    Retrieve indexed law metadata and article text.
    """

    try:

        vectorstore = get_vectorstore()

        results = vectorstore._collection.get(

            where={
                "law_id": law_id
            },

            include=[
                "documents",
                "metadatas",
            ],
        )

        if (
            not results
            or not results.get("documents")
        ):

            raise HTTPException(
                status_code=404,
                detail=(
                    f"No law found with id "
                    f"'{law_id}'."
                ),
            )

        article_groups: Dict[
            Any,
            list
        ] = {}

        law_title = "Legal Provision"

        year = ""

        category = ""

        source_file = ""

        for doc_text, meta in zip(

            results.get(
                "documents",
                [],
            ),

            results.get(
                "metadatas",
                [],
            ),

        ):

            meta = meta or {}

            article_number = meta.get(
                "article_number",
                0,
            )

            chunk_id = meta.get(
                "chunk_id",
                0,
            )

            clean = re.sub(
                r"^\[[^\]]+\]\s*",
                "",
                doc_text or "",
            )

            article_groups.setdefault(
                article_number,
                [],
            ).append(
                (
                    chunk_id,
                    clean,
                )
            )

            law_title = meta.get(
                "law_title",
                law_title,
            )

            year = meta.get(
                "year",
                year,
            )

            category = meta.get(
                "category",
                category,
            )

            source_file = meta.get(
                "source",
                source_file,
            )

        articles = []

        for article_number in sorted(
            article_groups,
            key=lambda value: (
                int(value)
                if str(value).isdigit()
                else 0
            ),
        ):

            chunks = sorted(
                article_groups[
                    article_number
                ],
                key=lambda x: x[0],
            )

            combined = re.sub(
                r"\s+",
                " ",
                " ".join(
                    text
                    for _, text in chunks
                ),
            ).strip()

            label = (
                f"Article {article_number}"
                if article_number
                else "General provisions"
            )

            articles.append(
                {
                    "number": label,
                    "text": combined[:3000],
                }
            )

        return {

            "law_id": law_id,

            "title": law_title,

            "full_text": (
                "\n\n".join(
                    f"{article['number']}\n"
                    f"{article['text']}"
                    for article in articles
                )
            )[:8000],

            "articles": articles[:30],

            "category": category,

            "year": year,

            "source_file": source_file,
        }

    except HTTPException:

        raise

    except Exception as exc:

        logger.error(
            "Law detail error '%s': %s",
            law_id,
            exc,
            exc_info=True,
        )

        raise HTTPException(
            status_code=500,
            detail="Could not retrieve law details.",
        )


# ═════════════════════════════════════════════
# CLEAR SESSION
# ═════════════════════════════════════════════

@app.delete(
    "/api/session/{session_id}"
)
async def clear_session(
    session_id: str
):

    if agent is None:

        raise HTTPException(
            status_code=503,
            detail="Assistant not ready.",
        )

    try:

        # Clear actual AI memory.
        agent.clear(session_id)

        # Clear API session metadata.
        sessions.pop(
            session_id,
            None,
        )

        return {
            "message": (
                f"Session '{session_id}' "
                "cleared successfully."
            )
        }

    except Exception as exc:

        logger.error(
            "Session clear error: %s",
            exc,
            exc_info=True,
        )

        raise HTTPException(
            status_code=500,
            detail="Could not clear session.",
        )


# ═════════════════════════════════════════════
# CLEAR SESSION FILE
# ═════════════════════════════════════════════

@app.delete(
    "/api/session/{session_id}/file"
)
async def clear_session_file(
    session_id: str
):
    """
    Clears file-related session state.

    The v8 upload flow normally stores extracted text
    on the frontend and sends it as file_context.
    Therefore there may be no server-side document to
    delete.

    This endpoint is kept so the frontend has a stable
    lifecycle API and can later switch to server-side
    document storage without changing its route.
    """

    # Currently remove any tracked file metadata that
    # explicitly references this session.

    to_delete = [
        file_id
        for file_id, data in uploaded_files.items()
        if data.get("session_id") == session_id
    ]

    for file_id in to_delete:
        uploaded_files.pop(
            file_id,
            None,
        )

    return {
        "message": (
            f"File context for session "
            f"'{session_id}' cleared."
        ),
        "deleted_files": len(to_delete),
    }


# ═════════════════════════════════════════════
# DEV / DEBUG — ACTIVE SESSIONS
# ═════════════════════════════════════════════

@app.get("/api/sessions")
async def list_sessions():

    cleanup_stale_sessions()

    now = time.time()

    return {

        session_id: {

            "created_ago_minutes": round(
                (
                    now
                    - data["created_at"]
                ) / 60,
                1,
            ),

            "last_used_ago_seconds": round(
                now
                - data["last_used"],
                1,
            ),

        }

        for session_id, data
        in sessions.items()
    }


# ═════════════════════════════════════════════
# ADMIN — INGEST
# ═════════════════════════════════════════════

@app.post("/api/ingest")
async def ingest_endpoint(
    rebuild: bool = False
):
    """
    Manually trigger legal document ingestion.

    IMPORTANT:
    Protect this endpoint with admin authentication
    before exposing it publicly.
    """

    try:

        if ingest_all_documents is not None:

            ingest_all_documents(
                rebuild=rebuild
            )

        else:

            from services.vector_store import ingest_gazettes

            ingest_gazettes()

        count = (
            get_vectorstore()
            ._collection
            .count()
        )

        return {

            "status": "ok",

            "chunk_count": count,

            "rebuild": rebuild,

        }

    except TypeError:

        # Compatibility with older force_ingest.py
        # versions that do not accept rebuild=.

        try:

            if ingest_all_documents is not None:

                ingest_all_documents()

            else:

                from services.vector_store import ingest_gazettes

                ingest_gazettes()

            count = (
                get_vectorstore()
                ._collection
                .count()
            )

            return {

                "status": "ok",

                "chunk_count": count,

                "rebuild": False,

            }

        except Exception as exc:

            logger.error(
                "Ingest error: %s",
                exc,
                exc_info=True,
            )

            raise HTTPException(
                status_code=500,
                detail="Legal document ingestion failed.",
            )

    except Exception as exc:

        logger.error(
            "Ingest error: %s",
            exc,
            exc_info=True,
        )

        raise HTTPException(
            status_code=500,
            detail="Legal document ingestion failed.",
        )


# ═════════════════════════════════════════════
# ROOT
# ═════════════════════════════════════════════

@app.get("/")
async def root():

    return {

        "name": "Rwanda Legal AI",

        "version": "8.0.0",

        "status": (
            "ready"
            if assistant_ready
            else "initializing"
        ),

        "documentation": "/docs",

    }


# ═════════════════════════════════════════════
# RUN
# ═════════════════════════════════════════════

if __name__ == "__main__":

    uvicorn.run(

        "main:app",

        host="0.0.0.0",

        port=8000,

        reload=True,

        log_level="info",
    )