import gradio as gr
import random
import string
from collections import defaultdict

# Store all conversations and users
conversations = defaultdict(list)
usernames = defaultdict(list)

# Generate a random conversation ID
def generate_conversation_id():
    return ''.join(random.choices(string.ascii_letters + string.digits, k=8))

# Handle new chat creation
def new_chat():
    conversation_id = generate_conversation_id()
    conversations[conversation_id] = []  # Initialize the conversation in the dictionary
    return conversation_id

# Handle user joining the conversation
def join_chat(chat_id, username):
    if username not in usernames[chat_id]:
        usernames[chat_id].append(username)
    return f"Welcome {username}! You have joined the chat."

# Handle sending and receiving messages
def chat_message_interface(chat_id, username, message):
    if chat_id not in conversations:
        return "Invalid chat ID. Please start a new chat."
    if username not in usernames.get(chat_id, []):
        return "You need to join the chat first by providing your username."
    
    # Store the message
    conversations[chat_id].append((username, message))
    
    # Display all messages
    all_messages = "\n".join([f"{user}: {msg}" for user, msg in conversations[chat_id]])
    return all_messages

# Gradio app
def main():
    with gr.Blocks() as demo:
        gr.Markdown("# Chat Application\nChoose an option to start or join a conversation.")
        
        # Initial buttons for New Chat or Enter Conversation ID
        with gr.Column() as initial_buttons:
            new_chat_button = gr.Button("New Chat")
            existing_chat_button = gr.Button("Enter Conversation ID for Existing Chat")
        
        # Chat interface components, initially hidden
        with gr.Column(visible=False) as chat_interface_column:
            chat_id_display = gr.Markdown("Chat ID: ")
            chat_id = gr.Textbox(label="Chat ID", placeholder="Enter Chat ID to join or create a new one.")
            username = gr.Textbox(label="Username", placeholder="Enter your username.")
            join_button = gr.Button("Join Chat")
            join_output = gr.Textbox(label="Join Status")
            message = gr.Textbox(label="Your Message", placeholder="Type your message here.")
            send_button = gr.Button("Send Message")
            chat_output = gr.Textbox(label="Chat Messages")
        
        # New Chat button logic
        def handle_new_chat():
            conversation_id = new_chat()
            return gr.update(visible=True), gr.update(value=f"Chat ID: {conversation_id}"), conversation_id
        
        new_chat_button.click(fn=handle_new_chat, outputs=[chat_interface_column, chat_id_display, chat_id])
        
        # Existing Chat button logic
        def handle_existing_chat():
            return gr.update(visible=True)
        
        existing_chat_button.click(fn=handle_existing_chat, outputs=chat_interface_column)
        
        # Join Chat button logic
        join_button.click(fn=join_chat, inputs=[chat_id, username], outputs=join_output)
        
        # Send Message button logic
        send_button.click(fn=chat_message_interface, inputs=[chat_id, username, message], outputs=chat_output)
    
    demo.launch()

if __name__ == "__main__":
    main()