import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Load the model and tokenizer once to avoid reloading
model_name = "moonshotai/Kimi-K2-Instruct"
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained(model_name)

def chat_with_kimi_k2(input_text, chat_history=None):
    if chat_history is None:
        chat_history = []
    # Append user input to chat history
    chat_history.append(f"User: {input_text}")
    
    # Prepare prompt
    prompt = "\n".join(chat_history) + "\nAI:"
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs, max_new_tokens=150, do_sample=True, temperature=0.7)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Extract AI response
    response_text = response.split("AI:")[-1].strip()
    chat_history.append(f"AI: {response_text}")
    
    return response_text, chat_history

# Main menu interface
def main_menu():
    return gr.Interface(
        fn=show_model_selection,
        title="AI Models Hub",
        description="Select an AI model to chat with:",
        inputs=[],
        outputs=gr.Button("Kimi K2"),
        live=True
    )

# Show model options
def show_model_selection():
    with gr.Row():
        kimi_button = gr.Button("Chat with Kimi K2")
    return kimi_button

# Chat interface for Kimi K2
def kimi_chat():
    chat_history = []

    def respond(user_input, history):
        response, chat_history_updated = chat_with_kimi_k2(user_input, chat_history=history)
        return response, chat_history_updated

    with gr.Blocks():
        gr.Markdown("### Chat with Kimi K2")
        chatbox = gr.Textbox(placeholder="Type your message...", label="You")
        response_box = gr.Textbox(label="Kimi K2")
        chat_history_state = gr.State([])
        back_button = gr.Button("Back to Menu")
        
        def back():
            return "Back", None

        back_button.click(fn=main_menu, inputs=None, outputs=None)
        
        user_input = chatbox
        user_input.submit(respond, [user_input, chat_history_state], [response_box, chat_history_state])
        back_button.click(lambda: None, None, None)

        return gr.Column([chatbox, response_box, back_button])

# Assemble the main app
def app():
    with gr.Blocks() as demo:
        menu = gr.Button("Go to Main Menu")
        kimichat = gr.Button("Chat with Kimi K2")
        
        def show_kimi():
            return kimi_chat()

        menu.click(show_kimi)
        kimichat.click(show_kimi)

    return demo

# Run the app
demo = app()
demo.launch()