from fastapi import FastAPI
from pydantic import BaseModel
import tensorflow as tf
import sentencepiece as spm
import numpy as np
import traceback

app = FastAPI()

# Load tokenizer and model
sp = spm.SentencePieceProcessor()
sp.load('sentencepiece.bpe.model')

interpreter = tf.lite.Interpreter(model_path="toxic_xlmr.tflite")
interpreter.allocate_tensors()

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
max_len = input_details[0]['shape'][1]


def run_inference(text: str) -> bool:
    """Run the model on a single text and return True if toxic."""
    tokens = [0] + sp.encode_as_ids(text) + [2]

    if len(tokens) > max_len:
        tokens = tokens[:max_len]
        tokens[-1] = 2
    else:
        tokens = tokens + [1] * (max_len - len(tokens))

    for i, detail in enumerate(input_details):
        shape = detail['shape']
        dtype = detail['dtype']
        input_data = np.zeros(shape, dtype=dtype)

        if shape[1] == max_len:
            if 'mask' in detail['name'].lower():
                mask = [1] * min(len(sp.encode_as_ids(text)) + 2, max_len)
                mask += [0] * (max_len - len(mask))
                input_data[0] = np.array(mask, dtype=dtype)
            elif 'type' in detail['name'].lower():
                pass
            else:
                input_data[0] = np.array(tokens, dtype=dtype)

        interpreter.set_tensor(detail['index'], input_data)

    interpreter.invoke()
    output_data = interpreter.get_tensor(output_details[0]['index'])

    if output_data.shape[-1] == 2:
        return bool(output_data[0][1] > 0.5)
    else:
        return bool(output_data[0][0] > 0.5)


class TextRequest(BaseModel):
    text: str


@app.post("/predict")
def predict(request: TextRequest):
    try:
        text = request.text

        # Test multiple casing variations against the model.
        # If ANY variation is toxic, we flag the whole text.
        variations = set()
        variations.add(text)                # Original: "Fuck you"
        variations.add(text.lower())        # Lowercase: "fuck you"
        variations.add(text.upper())        # Uppercase: "FUCK YOU"
        variations.add(text.capitalize())   # Capitalize: "Fuck you"
        variations.add(text.title())        # Title: "Fuck You"

        for variant in variations:
            if run_inference(variant):
                return {"is_toxic": True}

        return {"is_toxic": False}
    except Exception as e:
        error_msg = traceback.format_exc()
        print(error_msg)
        return {"is_toxic": False, "error": error_msg}


@app.get("/")
def read_root():
    return {"status": "Toxic API is running!"}