from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Dict
from utils import mask_pii, demask_pii
from models import classify_email


app = FastAPI()


class Entity(BaseModel):
    position: List[int]
    classification: str
    entity: str


class RequestBody(BaseModel):
    email_body: str


class ResponseBody(BaseModel):
    input_email_body: str
    list_of_masked_entities: List[Entity]
    masked_email: str
    category_of_the_email: str


@app.get('/')
def index():
    return "hello world"


@app.post("/classify-email", response_model=ResponseBody)
def classify_email_endpoint(req: RequestBody):
    text = req.email_body
    if not text:
        raise HTTPException(status_code=400, detail="Empty email body")

    masked_text, entities = mask_pii(text)
    category = classify_email(masked_text)

    return ResponseBody(
        input_email_body=text,
        list_of_masked_entities=entities,
        masked_email=masked_text,
        category_of_the_email=category,
    )