import sys
import time
import gradio as gr

from config import initialize, check_user

# Define the available options
stages = [
    "Opportunity Identification",
    "Refinement of Business Concept & Resource Acquisition",
    "Survival and Growth"
]

strengths_by_stage = {
    "Opportunity Identification": ["Curiosity", "Perspective", "Creativity", "Hope"],
    "Refinement of Business Concept & Resource Acquisition": ["Hope", "Bravery", "Honesty", "Love of Learning", "Social Intelligence", "Kindness"],
    "Survival and Growth": ["Persistence", "Humor", "Self-Regulation", "Zest", "Leadership", "Teamwork", "Fairness"]
}

ranks = ["High", "Medium", "Low"]

# Set default values
default_stage = stages[0]
default_strengths = strengths_by_stage[default_stage]
default_strength = default_strengths[0]

# Function to update the Strength dropdown when a stage is selected
def update_strengths(stage):
    strengths = strengths_by_stage.get(stage, [])
    return gr.update(choices=strengths, value=strengths[0] if strengths else None)

# Function to store the selected values and show the chat section
selected_stage = None
selected_strength = None
selected_rank = None

def get_question(rank):
    if rank=="High":
        return "How would you like to make use of this strength in setting up your business?"
    elif rank == "Medium":
        return "How do you plan to leverage this strength in shaping your business approach? Would you prefer to develop it further, seek support, collaborate with others, or take a different path?"
    elif rank == "Low":
        return "How would you like to go about (e.g., ignore, develop by…, ask for support to…, team up with others, reconsider choice to be entrepreneur)?"
    else:
        return None

def set_chat_config(stage, strength, rank):
    global selected_stage, selected_strength, selected_rank
    selected_stage = stage
    selected_strength = strength
    selected_rank = rank
    chat_title = f"### Chat Configuration: {stage} > {strength} > {rank}"

    return (        
        gr.update(visible=False),
        gr.update(visible=True),
        gr.update(value=chat_title),
    )

def reset_configuration():

    return (        
        gr.update(visible=True),
        gr.update(visible=False)
    )

llm_manager, mongo_db_manager = initialize()
if llm_manager is None or mongo_db_manager is None:
    print("Error: Failed to initialize configuration. Exiting application.", flush=True)
    sys.exit(1)

def reset_textbox():
    """Clears the textbox after sending a message."""
    return gr.update(value="")

def format_history(history):
    """Converts role-based history into the format expected by gr.Chatbot."""
    messages = []
    user_msg = None
    for entry in history:
        if entry["role"] == "user":
            user_msg = entry["content"]  # Store user message
        elif entry["role"] == "assistant" and user_msg:
            messages.append((user_msg, entry["content"]))  # Pair user with bot response
            user_msg = None  # Reset for next input
    return messages

def slow_echo(message, history):
    if history is None:
        history = []  # Ensure history is initialized

    # Append user message with role "user"
    history.append({"role": "user", "content": message})

    # Placeholder for assistant response
    bot_entry = {"role": "assistant", "content": ""}
    history.append(bot_entry)

    response = "You typed: "

    for i in range(len(message)):
        time.sleep(0.05)
        response += message[i]
        bot_entry["content"] = response  # Update assistant's response progressively
        yield history  # Yield updated history in the correct format

    yield history  # Final yield with full message

def llm_setup(stage, strength, rank, history):
    questions = mongo_db_manager.get_questions("questions", stage, strength, rank)
    if not questions:
        yield [{"role": "assistant", "content": "Error: no questions found for the selected configuration. Please try again."}]
        return

    question = get_question(rank)
    llm_manager.reset_messages(stage, strength, rank, question, questions)

    message = "Can you help me with this task?"
    yield from llm_send_message(message, history)

def llm_send_message(message, history):
    if history is None:
        history = []

    # Append user message to history
    history.append({"role": "user", "content": message})

    yield history 

    # Placeholder for assistant response
    bot_entry = {"role": "assistant", "content": ""}
    history.append(bot_entry)

    # Send message to LLM and stream response
    response = ""
    for chunk in llm_manager.send_message(message):  # Streaming response
        time.sleep(0.01)  # Simulate gradual output
        response += chunk
        bot_entry["content"] = response  # Update assistant response progressively
        yield history  # Yield updated history

    yield history  # Final yield

def authenticate(username, password):
    if check_user(username, password):
        print("🔑 Login successful!")
        return gr.update(visible=False), gr.update(visible=True), gr.update(value="", visible=False)  # Hide login, show chatbot, clear error
    else:
        print("❌ Incorrect username or password")
        return gr.update(visible=True), gr.update(visible=False), gr.update(value="❌ Incorrect username or password", visible=True)  # Show error

with gr.Blocks(fill_height=True) as demo:
    
    with gr.Column(visible=True) as login_section:
        gr.Markdown("### 🔒 Login Required")
        username_input = gr.Textbox(label="Username")
        password_input = gr.Textbox(label="Password", type="password")
        login_button = gr.Button("Login")
        error_message = gr.Text("", visible=False)

    with gr.Column(visible=False) as setup_section:
        gr.Markdown("### Select your Stage, Strength, and Rank before chatting")

        # Dropdowns for selection
        stage_dropdown = gr.Dropdown(
            label="Stage",
            choices=stages,
            value=default_stage,
            interactive=True
        )

        strength_dropdown = gr.Dropdown(
            label="Strength",
            choices=default_strengths,
            value=default_strength,
            interactive=True
        )

        rank_dropdown = gr.Dropdown(
            label="Rank",
            choices=ranks,
            value="High",
            interactive=True
        )

        # Update strengths when stage is changed
        stage_dropdown.change(
            fn=update_strengths,
            inputs=[stage_dropdown],
            outputs=[strength_dropdown]
        )

        # Button to set chat configuration
        setup_button = gr.Button("Set Chat Configuration")

    with gr.Column(visible=False) as chat_section:

        chat_configuration = gr.Markdown("")

        reset_button = gr.Button("Reset Chat Configuration")
    
        chat = gr.Chatbot(
            label="My strengths as an entrepreneur",
            type="messages"
        ) 

        input = gr.Textbox(
            label="Input", 
            placeholder="Type something here..."
        )

        stored_message = gr.State()

        input.submit(
            fn=lambda text: (text, ""),  # ✅ Store input in state, clear textbox immediately
            inputs=[input],  
            outputs=[stored_message, input]  # ✅ Store value & reset input
        ).then(
            fn=llm_send_message,  
            inputs=[stored_message, chat],  # ✅ Use stored message
            outputs=chat
        )

        send_btn = gr.Button("Send")  
        send_btn.click(
            fn=lambda text: (text, ""),  
            inputs=[input],  
            outputs=[stored_message, input]  
        ).then(
            fn=llm_send_message,  
            inputs=[stored_message, chat],  
            outputs=chat
        )

    login_button.click(
        authenticate, 
        [username_input, password_input], 
        [login_section, setup_section, error_message]
    )

    setup_button.click(
        fn=set_chat_config,
        inputs=[stage_dropdown, strength_dropdown, rank_dropdown],
        outputs=[setup_section, chat_section, chat_configuration]
    ).then(
        fn=llm_setup,
        inputs=[stage_dropdown, strength_dropdown, rank_dropdown, chat],
        outputs=chat
    )

    reset_button.click(
        fn=reset_configuration,
        inputs=None, 
        outputs=[setup_section, chat_section]
    )

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