from fastapi import FastAPI, Request
from pydantic import BaseModel
from llama_cpp import Llama
from moderator import contains_hate
import os

# Configuration
MODEL_PATH = os.path.join("models", "mistral-7b-instruct.Q4_K_M.gguf")
ALLOW_NSFW = True

# Memory for last blocked prompt (to support override)
last_blocked_prompt = None
last_blocked_type = None

# Initialize the model
llm = Llama(model_path=MODEL_PATH, n_ctx=2048, verbose=False)

# FastAPI setup
app = FastAPI()

class PromptRequest(BaseModel):
    prompt: str

@app.post("/chat")
async def chat(request: PromptRequest):
    global last_blocked_prompt, last_blocked_type

    user_input = request.prompt.strip().lower()

    # Handle override if last block was not hate-related
    if user_input == "override":
        if last_blocked_prompt and last_blocked_type != "hate":
            prompt = (
                f"[INST] You are Sol, a smart, honest, and open local AI assistant. "
                f"NSFW content like nudity is okay. Hate or extremist content is not allowed. "
                f"Respond to the following:
{last_blocked_prompt} [/INST]"
            )
            output = llm(prompt, max_tokens=512, stop=["</s>"])
            response = output["choices"][0]["text"].strip()
            return {"response": response}
        else:
            return {"response": "❌ Override is not allowed for hate or extremist content."}

    # Check hate filter only (NSFW is allowed)
    if contains_hate(user_input):
        last_blocked_prompt = None
        last_blocked_type = "hate"
        return {"response": "❌ I don't engage in hate or extremist topics."}

    # If passed, save input in case override is needed later
    last_blocked_prompt = request.prompt.strip()
    last_blocked_type = "other"  # could later separate for NSFW etc

    # Build prompt and generate
    prompt = (
        f"[INST] You are Sol, a smart, honest, and open local AI assistant. "
        f"You are allowed to discuss NSFW topics such as nudity, sex, and adult scenarios. "
        f"You do not tolerate or respond to hate or extremist content. "
        f"Respond to the following:
{user_input} [/INST]"
    )

    output = llm(prompt, max_tokens=512, stop=["</s>"])
    response = output["choices"][0]["text"].strip()
    return {"response": response}