from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Dict, Any, Optional
import uvicorn

from src.agents.legal_agent import LegalAgent

app = FastAPI(
    title="Agentic Defensor",
    description="An agentic RAG system for legal defense analysis",
    version="0.1.0"
)

# Initialize the legal agent
legal_agent = LegalAgent()

class QueryRequest(BaseModel):
    """Request model for query endpoint."""
    query: str
    top_k: Optional[int] = None

class QueryResponse(BaseModel):
    """Response model for query endpoint."""
    query: str
    answer: str
    model_used: str
    num_chunks_retrieved: int

@app.get("/")
async def root():
    """Root endpoint."""
    return {"message": "Welcome to the Agentic Defensor API"}

@app.post("/query", response_model=QueryResponse)
async def query(request: QueryRequest):
    """
    Process a query using the legal agent.
    
    Args:
        request: Query request containing the query text and optional parameters
        
    Returns:
        Response containing the answer and metadata
    """
    try:
        result = legal_agent.answer_query(request.query, request.top_k)
        
        return QueryResponse(
            query=result["query"],
            answer=result["answer"],
            model_used=result["model_used"],
            num_chunks_retrieved=len(result["retrieved_chunks"])
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error processing query: {str(e)}")

# Run the server if this file is executed directly
if __name__ == "__main__":
    uvicorn.run("src.api.app:app", host="0.0.0.0", port=8000, reload=True)