import os
import random
import gradio as gr
import sentencepiece as spm
import numpy as np
import pandas as pd
import tensorflow as tf
from valx import detect_profanity, detect_hate_speech

# Configuration dictionary mapping game names to their max_seq_len
MODEL_CONFIGS = {
    "Terraria": 12, "Skyrim": 13, "Witcher": 20, "WOW": 16, "Minecraft": 17,
    "Dark Souls": 13, "Final Fantasy": 14, "Elden Ring": 18, "Zelda": 15,
    "Dragon Age": 16, "Fallout": 13, "Darkest Dungeon": 14, "Monster Hunter": 15,
    "Bloodborne": 12, "Hollow Knight": 15, "Assassin's Creed": 15, "Baldur's Gate": 14,
    "Cyberpunk": 11, "Mass Effect": 13, "God Of War": 12, "Last Of Us": 5,
    "Factorio": 8, "The Sims": 4, "Fortnite": 9, "League Of Legends": 12,
    "Among Us": 13, "Warframe": 13, "Call of Duty": 11, "Forza Horizon": 10,
    "Halo": 14, "Overwatch": 9, "Subnautica": 14, "Fantasy": 16,
    
    # New Models
    "Animal Crossing": 14, "Civilization VI": 22, "Control": 22, "Cuphead": 24,
    "Dead Space": 18, "Diablo": 20, "Dota 2": 27, "EVE Online": 24, "GTA": 23,
    "Hades": 20, "Metroid": 28, "Portal": 28, "Resident Evil": 21, "RimWorld": 23,
    "Slay the Spire": 18, "Stardew Valley": 19, "Stellaris": 23, "Valheim": 20
}

# Global dictionary to store loaded models in memory
MODEL_CACHE = {}

def get_loaded_models(game_identifier):
    """
    Lazy-loads models into memory. If the model is already in the cache, 
    it returns it instantly. Otherwise, it loads it from disk, caches it, and returns it.
    """
    if game_identifier not in MODEL_CACHE:
        file_prefix = game_identifier.lower().replace(" ", "_").replace("'", "")
        
        # Load SentencePiece Model
        sp = spm.SentencePieceProcessor()
        sp.load(f"models/{file_prefix}_names.model")

        # Load TFLite model
        interpreter = tf.lite.Interpreter(model_path=f"models/dungen_{file_prefix}_model.tflite")
        interpreter.allocate_tensors()
        
        # Store in cache
        MODEL_CACHE[game_identifier] = {
            "sp": sp,
            "interpreter": interpreter,
            "vocab_size": sp.GetPieceSize()
        }
        
    return MODEL_CACHE[game_identifier]

def custom_pad_sequences(sequences, maxlen, padding='pre', value=0):
    padded_sequences = np.full((len(sequences), maxlen), value)
    for i, seq in enumerate(sequences):
        if padding == 'pre':
            if len(seq) <= maxlen:
                padded_sequences[i, -len(seq):] = seq
            else:
                padded_sequences[i, :] = seq[-maxlen:]
        elif padding == 'post':
            if len(seq) <= maxlen:
                padded_sequences[i, :len(seq)] = seq
            else:
                padded_sequences[i, :] = seq[:maxlen]
    return padded_sequences

def generate_random_name(interpreter, vocab_size, sp, max_length=10, temperature=0.5, seed_text="", max_seq_len=12):
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()
    decoded_name = ''

    if seed_text:
        generated_name = seed_text
    else:
        random_index = np.random.randint(1, vocab_size)
        random_token = sp.id_to_piece(random_index)
        generated_name = random_token

    for _ in range(max_length - 1):
        token_list = sp.encode_as_ids(generated_name)

        if len(token_list) == 0:
            continue  
            
        token_list = custom_pad_sequences([token_list], maxlen=max_seq_len, padding='pre') 
        token_list = token_list.astype(np.float32)

        interpreter.set_tensor(input_details[0]['index'], token_list)
        interpreter.invoke()

        predicted = interpreter.get_tensor(output_details[0]['index'])[0]
        predicted = np.log(predicted + 1e-8) / temperature
        predicted = np.exp(predicted) / np.sum(np.exp(predicted))

        next_index = np.random.choice(range(vocab_size), p=predicted)
        next_index = int(next_index)
        next_token = sp.id_to_piece(next_index)

        generated_name = sp.decode_pieces(sp.encode_as_pieces(generated_name) + [next_token])
        decoded_name = sp.decode_pieces(sp.encode_as_pieces(generated_name))

        if next_token == '' or len(decoded_name) > max_length:
            break

    # --- TEXT NORMALIZATION BLOCK ---
    
    # 1. Strip out unwanted tokens (including <s>)
    decoded_name = decoded_name.replace("▁", " ")
    decoded_name = decoded_name.replace("</s>", "")
    decoded_name = decoded_name.replace("<unk>", "")
    decoded_name = decoded_name.replace("<s>", "") 

    # 2 & 3. Normalize spacing and apply Title Case to every word
    words = decoded_name.split()
    normalized_name = " ".join([word.capitalize() for word in words])

    # 4. Split the name and check the last part length rule
    parts = normalized_name.split()
    if parts and len(parts[-1]) < 3:
        normalized_name = " ".join(parts[:-1])
        
    return normalized_name.strip()

# Note: Preserving the exact parameter names (like 'type') to ensure the API contract remains unbroken
def generateNames(type, amount, max_length=30, temperature=0.5, seed_text=""):
    hate_speech = detect_hate_speech(seed_text)
    profanity = detect_profanity([seed_text], language='All')
    
    if len(profanity) > 0:
        gr.Warning("Profanity detected in the seed text, using an empty seed text.")
        seed_text = ''
    else:
        if hate_speech == ['Hate Speech']:
            gr.Warning('Hate speech detected in the seed text, using an empty seed text.')
            seed_text = ''
        elif hate_speech == ['Offensive Speech']:
            gr.Warning('Offensive speech detected in the seed text, using an empty seed text.')
            seed_text = ''

    if type not in MODEL_CONFIGS:
        return pd.DataFrame([], columns=['Names'])

    # Fetch max sequence length
    max_seq_len = MODEL_CONFIGS[type]
    
    # Fetch cached models (loads them instantly if already cached)
    cached_data = get_loaded_models(type)
    sp = cached_data["sp"]
    interpreter = cached_data["interpreter"]
    vocab_size = cached_data["vocab_size"]

    amount = int(amount)
    max_length = int(max_length)
    names = []

    for _ in range(amount):
        generated_name = generate_random_name(
            interpreter, vocab_size, sp, 
            seed_text=seed_text, max_length=max_length, 
            temperature=temperature, max_seq_len=max_seq_len
        )
        stripped = generated_name.strip()
        
        # In case the generation completely fails and returns empty
        if not stripped:
            names.append("Generation Failed")
            continue
            
        item_hate_speech = detect_hate_speech(stripped)
        item_profanity = detect_profanity([stripped], language='All')
        name = ''

        if len(item_profanity) > 0:
            name = "Profanity Detected"
        elif item_hate_speech == ['Hate Speech']:
            name = 'Hate Speech Detected'
        elif item_hate_speech == ['Offensive Speech']:
            name = 'Offensive Speech Detected'
        else:
            # Catch-all: If it's safe (or even if valx returns an empty list), keep the name!
            name = stripped
        
        names.append(name)
        
    return pd.DataFrame(names, columns=['Names'])

demo = gr.Interface(
    fn=generateNames,
    inputs=[
        gr.Radio(
            choices=list(MODEL_CONFIGS.keys()), 
            label="Choose a model for your request", 
            value="Terraria"
        ), 
        gr.Slider(1, 100, step=1, label='Amount of Names', info='How many names to generate, must be greater than 0'), 
        gr.Slider(5, 60, value=30, step=1, label='Max Length', info='Max length of the generated word'), 
        gr.Slider(0.1, 1, value=0.5, label='Temperature', info='Controls randomness of generation, higher values = more creative, lower values = more probalistic'), 
        gr.Textbox('', label='Seed text (optional)', info='The starting text to begin with', max_lines=1)
    ],
    # Removed row_count and column_count to prevent the dataframe from aggressively paginating or limiting the array visualization
    outputs=[gr.Dataframe(label="Generated Names", headers=["Names"])],
    title='Dungen - Name Generator',
    description=(
        "A fun game-inspired name generator. For an example of how to create, and train your model, like this one, head over to: https://github.com/Infinitode/OPEN-ARC/tree/main/Project-5-TWNG. There you will find our base model, the dataset we used, and implementation code in the form of a Jupyter Notebook (exported from Kaggle).\n\n"
        "Try Dungen online: [Dungen AI | Advanced Neural Name Generator](https://infinitode.netlify.app/experiments/dungen-ai/)"
    )
)

if __name__ == "__main__":
    # Added ssr_mode=False to bypass the new Gradio 5 Async loop teardown bug
    demo.launch(ssr_mode=False)