from fastapi import FastAPI, UploadFile, Form, Body, Query
from typing import Optional
from classifier_service import classify_text, classify_pdf

app = FastAPI(
    title="Legal AI Multi-Task Classifier",
    description="""
Legal AI Multi-Task Classifier API  
Supports:
- **Text input** via form-data or JSON body  
- **File upload** (.txt or .pdf)  

Optional query parameters for dynamic filtering:
- **top_k**: Number of top ILSI/ISS sections to return (default 5)
- **threshold**: Minimum score for predictions (default 0.5)
""",
    version="1.0.0"
)


@app.post("/classify", summary="Classify legal text or file", tags=["Classification"])
async def classify(
    file: Optional[UploadFile] = None,
    text: Optional[str] = Form(None, description="Text input via form-data"),
    json_text: Optional[str] = Body(None, embed=True, description="Text input via JSON body"),
    top_k: int = Query(5, description="Number of top sections to return"),
    threshold: float = Query(0.5, description="Score threshold for predictions")
):
    """
    Classify text or file input with dynamic top_k and score_threshold.

    **Examples:**

    - **Text input (form-data)**:  
      `POST /classify?top_k=3&threshold=0.6` with `text="The accused violated sections 420 and 406 of IPC"`

    - **JSON input**:  
      ```json
      POST /classify?top_k=5&threshold=0.5
      {
        "json_text": "The accused violated sections 420 and 406 of IPC"
      }
      ```

    - **File upload**:  
      Upload `.txt` or `.pdf` file and include optional query params `top_k` and `threshold`.
    """
    if file:
        if file.filename.endswith(".txt"):
            text_content = (await file.read()).decode("utf-8")
            return classify_text(text_content, top_k=top_k, threshold=threshold)
        elif file.filename.endswith(".pdf"):
            file_bytes = await file.read()
            return classify_pdf(file_bytes, top_k=top_k, threshold=threshold)
        else:
            return {"error": "Only .txt or .pdf files are supported."}

    final_text = text or json_text
    if final_text:
        return classify_text(final_text, top_k=top_k, threshold=threshold)

    return {"error": "Provide text input or upload a file."}


@app.get("/", summary="Home endpoint", tags=["General"])
def home():
    """
    Returns a simple message indicating the API is running.
    """
    return {"message": "Legal AI Multi-Task Classifier API"}