from fastapi import FastAPI, HTTPException
from typing import Any, Dict, List, Optional
from pydantic import BaseModel
from os import getenv
from huggingface_hub import InferenceClient
import random
import nltk
import re
from word_forms.word_forms import get_word_forms

app = FastAPI()

nltk.download('punkt')

tokenizer = nltk.data.load('tokenizers/punkt/english.pickle')

HF_TOKEN = getenv("HF_TOKEN")

class InputData(BaseModel):
    model: str
    system_prompt_template: List[str]
    prompt_template: List[str]
    end_token: str
    system_prompt: List[str]
    user_input: List[str]
    history: str = ""
    segment: bool = False
    max_sentences: Optional[int] = None

class WordCheckData(BaseModel):
    string: str
    word: str

@app.post("/generate-response/")
async def generate_response(data: InputData) -> Dict[str, Any]:
    if data.max_sentences is not None and data.max_sentences != 0:
        data.segment = True
    elif data.max_sentences == 0:
        for prompt in data.prompt_template:
            for user_input in data.user_input:
                data.history += prompt.replace("{Prompt}", user_input) + "\n"
        return {
            "response": "",
            "history": data.history + data.end_token
        }

    user_input_str = ""
    if data.segment:
        for user_input in data.user_input:
            user_sentences = tokenizer.tokenize(user_input)
            user_input_str += "\n".join(user_sentences) + "\n"
    else:
        user_input_str = "\n".join(data.user_input)

    for prompt in data.prompt_template:
        data.history += prompt.replace("{Prompt}", user_input_str) + "\n"

    inputs = ""
    for system_prompt in data.system_prompt_template:
        inputs += system_prompt.replace("{SystemPrompt}", "\n".join(data.system_prompt)) + "\n"
    inputs += data.history

    seed = random.randint(0, 2**32 - 1)

    try:
        client = InferenceClient(model=data.model, token=HF_TOKEN)
        response = client.text_generation(
            inputs,
            temperature=1.0,
            max_new_tokens=1000,
            seed=seed
        )

        response_str = str(response)

        if data.segment:
            ai_sentences = tokenizer.tokenize(response_str)
            if data.max_sentences is not None:
                ai_sentences = ai_sentences[:data.max_sentences]
            ai_response_str = "\n".join(ai_sentences)
        else:
            ai_response_str = response_str

        data.history += ai_response_str + "\n"

        cleaned_response = {
            "New response": ai_sentences if data.segment else [response_str],
            "Sentence count": len(ai_sentences) if data.segment else 1
        }

        return {
            "response": cleaned_response,
            "history": data.history + data.end_token
        }

    except Exception as e:
        print(f"Model {data.model} failed with error: {e}")
        raise HTTPException(status_code=500, detail=f"Model {data.model} failed to generate response")

@app.post("/check-word/")
async def check_word(data: WordCheckData) -> Dict[str, Any]:
    input_string = data.string.lower()
    word = data.word.lower()

    forms = get_word_forms(word)

    all_forms = set()
    for words in forms.values():
        all_forms.update(words)

    words_in_string = re.findall(r'\b\w+\b', input_string)

    found = any(word_in_string in all_forms for word_in_string in words_in_string)

    result = {
        "found": found
    }

    return result