import gradio as gr
from huggingface_hub import InferenceClient
from openai import OpenAI
import anthropic
import os
import json
from datetime import datetime

client = InferenceClient()

# Initialize OpenAI client
openai_client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

# Initialize Anthropic client
anthropic_client = anthropic.Anthropic(api_key=os.getenv("CLAUDE_API_KEY"))

MODEL_OPTIONS = {
    "Meta-Llama-3-8B-Instruct": "meta-llama/Meta-Llama-3-8B-Instruct",
    "Mixtral-8x7B-Instruct-v0.1": "mistralai/Mixtral-8x7B-Instruct-v0.1",
    "Mixtral-8x22B-v0.1": "mistralai/Mixtral-8x22B-v0.1",
    "OpenAI GPT-3.5": "gpt-3.5-turbo",
    "OpenAI GPT-4": "gpt-4",
    "claude-3-5-sonnet": "claude-3-5-sonnet-20240320"
}

def generate_text(message, history, model_choice):

    if model_choice in ["OpenAI GPT-3.5", "OpenAI GPT-4"]:
            # OpenAI API call
            messages = [{"role": "system", "content": "You are an AI assistant. Respond to the user's questions accurately and concisely."}]
            for h in history:
                if isinstance(h, list) and len(h) == 2:
                    messages.append({"role": "user", "content": h[0]})
                    if h[1]:
                        messages.append({"role": "assistant", "content": h[1]})
            messages.append({"role": "user", "content": message})
            
            stream = openai_client.chat.completions.create(
                model=MODEL_OPTIONS[model_choice],
                messages=messages,
                stream=True
            )
            for chunk in stream:
                if chunk.choices[0].delta.content is not None:
                    yield chunk.choices[0].delta.content


    elif model_choice == "claude-3-5-sonnet-20240620":
        # Claude API call
        messages = [{"role": "system", "content": "You are an AI assistant. Respond to the user's questions accurately and concisely."}]
        for h in history:
            if isinstance(h, list) and len(h) == 2:
                messages.append({"role": "user", "content": h[0]})
                if h[1]:
                    messages.append({"role": "assistant", "content": h[1]})
        messages.append({"role": "user", "content": message})
        
        stream = anthropic_client.messages.stream(
            model=MODEL_OPTIONS[model_choice],
            messages=messages,
            max_tokens=1000
        )
        for chunk in stream:
            if chunk.delta.text is not None:
                yield chunk.delta.text

    else:
       model = MODEL_OPTIONS[model_choice]
        conversation = "You are an AI assistant. Respond to the user's questions accurately and concisely.\n\n"
        for h in history:
            if isinstance(h, list):
                if len(h) >= 1:
                    conversation += f"Human: {h[0]}\n"
                if len(h) >= 2 and h[1]:
                    conversation += f"AI: {h[1]}\n"
            elif isinstance(h, str):
                conversation += f"{h}\n"
            conversation += "\n"
        
        conversation += f"Human: {message}\nAI:"
        
        output = client.text_generation(
            conversation,
            model=model,
            max_new_tokens=1000,
            temperature=0.8,
            stream=True
        )
        
        for chunk in output:
            yield chunk


def save_dialogue(history):
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"dialogue_{timestamp}.json"
    
    dialogue_dict = {"conversation": []}
    for human, ai in history:
        dialogue_dict["conversation"].append({"human": human, "ai": ai})
    
    with open(filename, 'w') as f:
        json.dump(dialogue_dict, f, indent=2)
    
    return f"Dialogue saved to {filename}"

with gr.Blocks() as iface:
    gr.Markdown("# Chat with LLM Models")
    gr.Markdown("Select a model and start chatting!")
    
    model_dropdown = gr.Dropdown(
        choices=list(MODEL_OPTIONS.keys()),
        value="Meta-Llama-3-8B-Instruct",
        label="Select Model"
    )
    
    chatbot = gr.Chatbot()
    msg = gr.Textbox()
    clear = gr.Button("Clear")
    save = gr.Button("Save Dialogue")

    def user(user_message, history):
        return "", history + [[user_message, None]]

    def bot(history, model_choice):
        user_message = history[-1][0]
        bot_message = generate_text(user_message, history[:-1], model_choice)
        history[-1][1] = ""
        for chunk in bot_message:
            history[-1][1] += chunk
            yield history

    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, [chatbot, model_dropdown], chatbot
    )
    clear.click(lambda: None, None, chatbot, queue=False)
    save.click(save_dialogue, inputs=[chatbot], outputs=[gr.Textbox(label="Save Status")])

iface.launch(share=True)