import functools
import os
import uuid
import asyncio
import gradio as gr

from langchain_mcp_adapters.client import MultiServerMCPClient
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import ToolNode
from langgraph.graph import MessagesState, END, StateGraph
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
from langgraph.checkpoint.memory import MemorySaver



# Apply a theme
theme = gr.themes.Soft(
    primary_hue="blue", secondary_hue="sky", neutral_hue="slate"
).set(
    body_background_fill="linear-gradient(to right, #f0f4f8, #e6e9f0)",  # Light gradient background
    block_background_fill="white",
    block_border_width="1px",
    block_shadow="*shadow_drop_lg",
    button_primary_background_fill="*primary_500",
    button_primary_text_color="white",
    button_secondary_background_fill="*secondary_500",
    button_secondary_text_color="white",
)

with gr.Blocks(theme=theme, title="PMCP - Agentic Project Management") as demo:
    gr.Markdown(
        """
        # PMCP - Agentic Project Management
        Manage your projects with AI assistance - GitHub, Trello.
        """
    )
    
    with gr.Row():
        # Left column for environment variables
        with gr.Column(scale=1):
            gr.Markdown("### Environment Variables")
            github_repo = gr.Textbox(label="GitHub Repo", placeholder="username/repository")
            github_token = gr.Textbox(label="GitHub Token", placeholder="ghp_xxxxxxxxxxxx", type="password")
            trello_api = gr.Textbox(label="Trello API", placeholder="Your Trello API key")
            trello_token = gr.Textbox(label="Trello Token", placeholder="Your Trello token", type="password")
            hf_token = gr.Textbox(label="HF Token", placeholder="Your Hugging Face token", type="password")
            
            set_env_button = gr.Button("Set Environment Variables", variant="primary")
            env_status = gr.Markdown("")
            
        # Right column for chat interface
        with gr.Column(scale=2):
            chatbot = gr.Chatbot(
                label="Conversation",
                bubble_full_width=False,
                height=600,
                avatar_images=(
                    None,
                    "https://upload.wikimedia.org/wikipedia/commons/thumb/7/72/Brandon_Sanderson_in_2018.jpg/800px-Brandon_Sanderson_in_2018.jpg?20230101015657",
                ),
                # (user_avatar, bot_avatar) - replace bot avatar with something better or remove
            )

            with gr.Row():
                msg = gr.Textbox(
                    label="Your Message",
                    placeholder="Type your message here and press Enter to send...",
                    scale=4,  # Make textbox take more space
                    autofocus=True,
                )
                with gr.Column(scale=1):
                    submit_button = gr.Button("Send", variant="primary")
                    clear_button = gr.Button("Reset Conversation", variant="secondary")

    async def respond(user_message_text, chat_history):
        if not user_message_text.strip():
            return (
                chat_history,
                "",
            )  # Ignore empty input, return current history and clear textbox

        is_first_turn = not chat_history

        # Append user message to chat_history optimistically for immediate display
        # Bot response will fill in the 'None' later or add new [None, bot_msg] rows
        # This makes UI feel more responsive.
        # chat_history.append([user_message_text, None]) # Temporarily removed for simpler logic below

        ai_utterances = await agent.generate_responses_for_turn(
            user_message_text, is_first_turn
        )

        if ai_utterances:
            chat_history.append([user_message_text, ai_utterances[0]])
            for i in range(1, len(ai_utterances)):
                chat_history.append([None, ai_utterances[i]])
        else:
            # If agent provides no utterances (e.g. error or no action)
            chat_history.append(
                [user_message_text, "I don't have a response for that right now."]
            )

        return chat_history, ""  # Return updated history and clear the textbox

    # Event handlers
    msg.submit(respond, [msg, chatbot], [chatbot, msg])
    submit_button.click(respond, [msg, chatbot], [chatbot, msg])

    def clear_chat_and_reset_agent():
        agent.reset_thread()
        return [], ""  # Clears chatbot UI and textbox
    
    def set_environment_variables(github_repo, github_token, trello_api, trello_token, hf_token):
        # Set environment variables
        if github_repo:
            os.environ["GITHUB_REPO"] = github_repo
        if github_token:
            os.environ["GITHUB_TOKEN"] = github_token
        if trello_api:
            os.environ["TRELLO_API_KEY"] = trello_api
        if trello_token:
            os.environ["TRELLO_TOKEN"] = trello_token
        if hf_token:
            os.environ["NEBIUS_API_KEY"] = hf_token
            
        # Create a message showing which variables were set
        set_vars = []
        if github_repo: set_vars.append("GITHUB_REPO")
        if github_token: set_vars.append("GITHUB_TOKEN")
        if trello_api: set_vars.append("TRELLO_API_KEY")
        if trello_token: set_vars.append("TRELLO_TOKEN")
        if hf_token: set_vars.append("NEBIUS_API_KEY")
        
        if set_vars:
            return f"✅ Set environment variables: {', '.join(set_vars)}"
        else:
            return "⚠️ No environment variables were set"

    # Connect the set environment variables button
    set_env_button.click(
        set_environment_variables,
        inputs=[github_repo, github_token, trello_api, trello_token, hf_token],
        outputs=[env_status]
    )
    
    clear_button.click(clear_chat_and_reset_agent, None, [chatbot, msg], queue=False)

    # Load agent setup when the app starts
    # Using a lambda to ensure asyncio.run is called within the demo's event loop context if needed
    demo.load(lambda: asyncio.run(agent.setup()))

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