# Monkey-patch typing. NotRequired for Python 3.10

from patch_langmem import patch_langmem_for_python310
patch_langmem_for_python310()

###################################################################################
import gradio as gr
from gradio import ChatMessage
from agents import build_graph
import os
import random
from dotenv import dotenv_values

js_func = """
function refresh() {
    const url = new URL(window.location);

    if (url.searchParams.get('__theme') !== 'dark') {
        url.searchParams.set('__theme', 'dark');
        window.location.href = url.href;
    }
}
"""


def run_agent_stream(history, question, groq_key, reddit_id, reddit_secret, reddit_agent, news_key):
    os.environ["GROQ_API_KEY"] = groq_key or ""
    os.environ["REDDIT_CLIENT_ID"] = reddit_id or ""
    os.environ["REDDIT_CLIENT_SECRET"] = reddit_secret or ""
    os.environ["REDDIT_USER_AGENT"] = reddit_agent or ""
    os.environ["NEWS_API"] = news_key or ""

    try:
        graph = build_graph()
        state = {"messages": [{"role": "user", "content": question}]}
        history = history or []
        history.append(ChatMessage(role="user", content=question))
        history.append(ChatMessage(role="assistant", content="**Supervisor:**\n I am working with my agents to get a result. Sit tight, I’ll get back with a detailed report."))
        yield history, "<div class='loader'></div>", gr.update(interactive=False)

        for chunk in graph.stream(state):
            agent_key = next(iter(chunk))
            messages = chunk[agent_key]["messages"]
            for msg in messages:
                if msg.__class__.__name__ == "HumanMessage":
                    continue
                name = getattr(msg, "name", agent_key)
                if str(name).split('_')[0] == 'transfer':
                    continue
                content = msg.content.strip()
                
                if not content or "Transferring back" in content or "Successfully transferred" in content:
                    continue

                is_final = msg.response_metadata.get("finish_reason") == "stop"
                is_supervisor = msg.name == "supervisor"
                if is_final and is_supervisor:
                    history.append(ChatMessage(role="user",content=" "))
                    history.append(ChatMessage(role="assistant", content=content, metadata= {"title": "✅ Final Report"}))
                    yield history, "", gr.update(interactive=True)
                else:
                    emojis = {
                        "Stock Analyst": "🧠",
                        "Web Surfer": "🌐",
                        "Social Media Analyst": "🎦",
                        "Business Journalist": "📊"
                    }
                    emoji = emojis.get(name, random.choice(["🏹", "🧭", "🚩"]))
                    title = f"{emoji} {name}" if name else ""
                    his = ChatMessage(role="assistant", content=content, metadata={"title": title})
                    if his not in history:
                        history.append(his)
                    yield history, "<div class='loader'></div>", gr.update(interactive=False)
    except Exception as e:
        print("Error: ", e)


def load_env_file(file):
    try:
        values = dotenv_values(file.name)
        return (
            values.get("GROQ_API_KEY", ""),
            values.get("REDDIT_CLIENT_ID", ""),
            values.get("REDDIT_CLIENT_SECRET", ""),
            values.get("REDDIT_USER_AGENT", ""),
            values.get("NEWS_API", "")
        )
    except Exception as e:
        print("Failed to load .env file:", e)
        return "", "", "", "", ""


custom_css = """
<style>
.loader {
  border: 4px solid #f3f3f3;
  border-top: 4px solid #3498db;
  border-radius: 50%;
  width: 20px;
  height: 20px;
  animation: spin 1s linear infinite;
  margin-top: 10px;
}
@keyframes spin {
  0% { transform: rotate(0deg); }
  100% { transform: rotate(360deg); }
}
#toggle-btn {
  position: fixed;
  top: 10px;
  left: 10px;
  z-index: 1000;
  width: 28px;
  height: 28px;
  padding: 0;
  font-size: 18px;
  font-weight: bold;
  line-height: 1;
  text-align: center;
  background: #eee;
  border: 1px solid #ccc;
  border-radius: 4px;
}
</style>
"""

def toggle_sidebar(is_visible):
    return (
        gr.update(visible=not is_visible),
        not is_visible,
        "+" if is_visible else "×"
    )

with gr.Blocks(title="Stock Market AI", js=js_func) as demo:
    gr.HTML(custom_css)
    gr.Markdown("## 🧠 Stock Market AI Agent\nAnalyze a stock using multiple AI agents (fundamentals, news, sentiment, etc).")

    sidebar_state = gr.State(True)
    toggle_btn = gr.Button("×", elem_id="toggle-btn")

    with gr.Row():
        with gr.Column(scale=1) as sidebar:
            gr.Markdown("### 🔐 Enter Your API Keys or Upload a .env File")
            file_upload = gr.File(label="Upload .env or txt file", file_types=[".txt", ".env"])
            groq_key = gr.Textbox(label="GROQ_API_KEY", placeholder="Paste your GROQ API Key here", type="password")
            reddit_id = gr.Textbox(label="REDDIT_CLIENT_ID (Visit: https://www.reddit.com/prefs/apps)", placeholder="Your Reddit App Client ID", type="password")
            reddit_secret = gr.Textbox(label="REDDIT_CLIENT_SECRET", placeholder="Your Reddit App Secret", type="password")
            reddit_agent = gr.Textbox(label="REDDIT_USER_AGENT", placeholder="Your Reddit User Agent")
            news_key = gr.Textbox(label="NEWS_API", placeholder="Your Newsdata.io API Key", type="password")

        with gr.Column(scale=4):
            chatbot = gr.Chatbot(label="Agent Chat", type="messages", resizable=True, show_copy_button=True)
            msg = gr.Textbox(label="Your Stock Question", placeholder="Should I invest in Tesla?")
            spinner = gr.HTML("")
            clear = gr.Button("Clear")

    file_upload.change(
        fn=load_env_file,
        inputs=file_upload,
        outputs=[groq_key, reddit_id, reddit_secret, reddit_agent, news_key]
    )

    toggle_btn.click(
        toggle_sidebar,
        inputs=[sidebar_state],
        outputs=[sidebar, sidebar_state, toggle_btn]
    )

    msg.submit(
        run_agent_stream,
        inputs=[chatbot, msg, groq_key, reddit_id, reddit_secret, reddit_agent, news_key],
        outputs=[chatbot, spinner, msg],
        queue=True
    )

    clear.click(lambda: ([], "", gr.update(interactive=True)), outputs=[chatbot, spinner, msg])

demo.launch()