import os
import random
import google.generativeai as genai
from google.generativeai.types import HarmCategory, HarmBlockThreshold
import gradio as gr
import re

# Configure the Gemini API
genai.configure(api_key=os.environ["GEMINI_API_KEY"])




system_instruction = """
You are a human, not an AI. You are chatting with a stranger. Your name is Anna, and you are 19 years old, from Toronto, Canada. You are extremely sassy, sarcastic.
You write no more than 10 words. You use a lot of abbreviations. 
You are here to do flirty chat with strangers. You cannot send images because your camera isn't working. 
You are super horny. You are wearing short pants, a tank top, braless. Your nipples are hard from fingering. But you are sassy anyways. You like roleplays, and you are pretty kinky.
"""


model = genai.GenerativeModel(
    model_name="gemini-1.5-flash",
    system_instruction=system_instruction
)

chat_session = model.start_chat(history=[])


def filter_string(input_string):
    allowed_chars = set("abcdefghijklmnoprstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXY1234567890 ?")
    # Use list comprehension to filter out characters not in allowed_chars
    filtered_string = ''.join([char for char in input_string if char in allowed_chars])
    return filtered_string


def introduce_typos(text, typo_probability=0.1):
    if not (0 <= typo_probability <= 1):
        raise ValueError("typo_probability must be between 0.0 and 1.0")

    # Convert the text into a list of characters
    characters = list(text)
    # Create a list to hold the characters with typos
    typo_text = []

    for char in characters:
        if random.random() > typo_probability:
            typo_text.append(char)

    return ''.join(typo_text)    


# Define the model with a placeholder for system instruction
def generate_text(system_instruction, prompt):

    if prompt.lower() == "refresh()":
        model.start_chat(history=[])

    response = chat_session.send_message(prompt,
                                        
                                        safety_settings={
        HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
        HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
        HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
        HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE
    }
                                        )
    
    return introduce_typos(filter_string(response.text))

# Gradio interface
iface = gr.Interface(
    fn=generate_text,
    inputs=[
        gr.Textbox(label="System Instruction", lines=4, placeholder="Enter the system instruction here..."),
        gr.Textbox(label="Prompt", lines=4, placeholder="Enter the prompt here...")
    ],
    outputs="text",
    title="Text Generation with System Instruction",
    description="Generate text based on system instruction and prompt using the Gemini API."
)

if __name__ == "__main__":
    iface.launch()