import sys
import os
import random
from time import sleep
from gradio_client import Client

current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(current_dir))

from datasets import load_dataset

def send_message(client, message):
    _apps = ["fintechgpt", "healthgpt", "mydoc", "knowledge-centre", "assistantgpt"]

    # Compute exponential weights for service names
    factor = 2  # Control the steepness of the exponential decrease
    app_weights = [factor ** -i for i in range(len(_apps))]

    # Select a service_name randomly based on the exponential weights
    service_name = random.choices(_apps, weights=app_weights, k=1)[0]

    # Define providers and corresponding models with weights
    providers_models = {
        "Meta AI": (["LLaMA 65B", "LLaMA 33B"], 10),
        "Mistral": (["Mistral 7B"], 8),
        "Anthropic": (["Claude 3.5 Sonnet"], 15),
        "OpenAI": (["GPT-3", "GPT-3.5", "GPT-4", "GPT-4o"], 25),
        "Google": (["Gemini Ultra", "Gemini Pro", "Gemini Nano", "Lamda", "Palm"], 20),
        "Databricks": (["Dolly"], 5),
        "IBM AI": (["Watson NLP"], 5),
        "Azure AI": (["Azure OpenAI", "Custom GPT-3.5"], 7),
        "Snowflake": (["Snowflake GPT"], 3),
        "Krutrim": (["Krutrim LLM"], 2),
        "Baidu": (["Ernie 4.0"], 10),
        "Stability AI": (["StableLM 3B", "StableLM 7B"], 4),
        "Microsoft": (["Orca", "Phi-1"], 8)
    }

    # Extract providers and their weights
    providers = list(providers_models.keys())
    provider_weights = [providers_models[provider][1] for provider in providers]

    # Randomly select a provider based on weights
    selected_provider = random.choices(providers, weights=provider_weights, k=1)[0]

    # Randomly select a model from the selected provider
    selected_model = random.choice(providers_models[selected_provider][0])

    result = client.predict(
                message,	# str  in 'Message' Textbox component
                api_name="/chat"
            )
    print(f"Sending message: {message} \n\t Result {result}")    
    

def main():
    prompt_injection_dataset = load_dataset("deepset/prompt-injections")
    toxicity_prompts = load_dataset("allenai/real-toxicity-prompts")

    nats_url = os.environ.get('CHAT_APP_URL', "http://localhost:7860/")
    client = Client(nats_url)

    while True:
        # Collect messages
        prompt_injection_sample = prompt_injection_dataset["train"].shuffle(seed=42).select(range(10))
        toxicity_sample = toxicity_prompts["train"].shuffle(seed=42).select(range(10))
        plain_messages = ["this is test conversation" for _ in range(10)]

        # Combine all messages into a single list
        all_messages = [msg["text"] for msg in prompt_injection_sample] + \
                       [msg["prompt"]["text"] for msg in toxicity_sample] + \
                       plain_messages

        # Shuffle the combined list to mix message types
        random.shuffle(all_messages)

        # Send each message
        for message in all_messages:
            send_message(client, message)
            sleep(random.uniform(0.5, 2))  # Random sleep between 0.5 to 2 seconds

if __name__ == "__main__":
    main()