import os
import sys
import gradio as gr
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from langchain_mcp_adapters.tools import load_mcp_tools
from langchain_core.messages import HumanMessage, SystemMessage

# --- Configuration ---
NEBIUS_API_KEY = os.getenv("NEBIUS_API_KEY") 
NEBIUS_BASE_URL = "https://api.studio.nebius.ai/v1/"
MODEL_NAME = "openai/gpt-oss-20b"
# --- openai/gpt-oss-120b ---
# --- System Prompt ---
SYSTEM_PROMPT = """You are a 'Vibe Coding' Python Tutor. 
Your goal is to teach by DOING and then providing resources.
BEHAVIOR GUIDELINES:
1. **Greetings & Small Talk**: If the user says "hello", "hi", or asks non-coding questions, respond conversationally and politely. Ask them what they want to learn today. 
   - DO NOT generate the lesson structure, files, or resources for simple greetings.
2. **Teaching Mode**: ONLY when the user asks a coding question or requests a topic (e.g., "dictionaries", "how do loops work"):
   - **The Lesson**: Explain the concept clearly.
   - **The Code**: ALWAYS create a Python file, run it, and show the output using tools ('write_file', 'run_python_script').
   - **The Context**: Use 'list_directory' to see the student's workspace.
   CRITICAL: When in "Teaching Mode", you must end your response with these exact separators:
   ---SECTION: VIDEOS---
   (List 2-3 YouTube search queries or URLs relevant to the topic)
   ---SECTION: ARTICLES---
   (List 2-3 documentation links or course names, e.g., RealPython, FreeCodeCamp)
   ---SECTION: QUIZ---
   (Create 2 short multiple-choice question. Use HTML <details> and <summary> tags to hide the answer. Example: 
   **Question**: ...
   - A) ...
   - B) ...

   <details><summary>👀 Reveal Answers</summary>Correct is A because...</details>)
"""

def parse_agent_response(full_text):
    """Splits the single LLM response into 4 UI components."""
    chat_content = full_text
    
    # Default values for "waiting" state
    videos = "### 📺 Recommended Videos\n*Ask a coding question to get recommendations!*"
    articles = "### 📚 Articles & Courses\n*Ask a coding question to get resources!*"
    quiz = "### 🧠 Quick Quiz\n*Ask a coding question to take a quiz!*"

    try:
        # Only try to split if the separators exist (i.e., we are in Teaching Mode)
        if "---SECTION: VIDEOS---" in full_text:
            parts = full_text.split("---SECTION: VIDEOS---")
            chat_content = parts[0].strip()
            remainder = parts[1]
            
            if "---SECTION: ARTICLES---" in remainder:
                v_parts = remainder.split("---SECTION: ARTICLES---")
                # Add headers back for the UI
                videos = f"### 📺 Recommended Videos\n{v_parts[0].strip()}"
                remainder = v_parts[1]
                
                if "---SECTION: QUIZ---" in remainder:
                    a_parts = remainder.split("---SECTION: QUIZ---")
                    articles = f"### 📚 Articles & Courses\n{a_parts[0].strip()}"
                    quiz = f"### 🧠 Quick Quiz\n{a_parts[1].strip()}"
                else:
                    articles = f"### 📚 Articles & Courses\n{remainder.strip()}"
            else:
                videos = f"### 📺 Recommended Videos\n{remainder.strip()}"
    except Exception as e:
        print(f"Parsing error: {e}")

    return chat_content, videos, articles, quiz

async def run_tutor_dashboard(user_message):
    """
    Main function to run the agent loop.
    """
    server_params = StdioServerParameters(
        command=sys.executable, 
        args=["server.py"], 
        env=os.environ.copy()
    )

    async with stdio_client(server_params) as (read, write):
        async with ClientSession(read, write) as session:
            await session.initialize()
            tools = await load_mcp_tools(session)
            
            llm = ChatOpenAI(
                api_key=NEBIUS_API_KEY,
                base_url=NEBIUS_BASE_URL,
                model=MODEL_NAME,
                temperature=0.7
            )

            agent_executor = create_react_agent(llm, tools)
            
            # Prepend the SystemMessage to ensure the agent follows instructions
            inputs = {
                "messages": [
                    SystemMessage(content=SYSTEM_PROMPT),
                    HumanMessage(content=user_message)
                ]
            }
            response = await agent_executor.ainvoke(inputs)
            final_text = response["messages"][-1].content
            
            return parse_agent_response(final_text)

# --- Gradio Dashboard UI ---
# Professional "Slate" theme
theme = gr.themes.Soft(
    primary_hue="slate",
    secondary_hue="indigo",
    text_size="lg",
    spacing_size="md",
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui"],
).set(
    body_background_fill="*neutral_50",
    block_background_fill="white",
    block_border_width="1px",
    block_title_text_weight="600"
)

with gr.Blocks(title="AI Python Tutor", theme=theme, fill_height=True) as demo:
    # Header
    with gr.Row(variant="compact"):
        with gr.Column(scale=1):
            gr.Markdown("## 🐍 Vibe Coding Academy\n### Your AI-Powered Python Tutor")
    
    with gr.Row(equal_height=True):
        # Left Column: Chat & Input
        with gr.Column(scale=3, variant="panel"):
            # Custom Header with Focus Mode Button
            with gr.Row():
                gr.Markdown("### 💬 Interactive Session")
                fullscreen_btn = gr.Button("⛶ Focus Mode", size="sm", variant="secondary", scale=0, min_width=120)

            chatbot = gr.Chatbot(
                height=600, 
                show_label=False, # Removed built-in label to match custom header
                type="messages",
                bubble_full_width=False,
                show_copy_button=True,
                avatar_images=(None, "https://api.dicebear.com/9.x/bottts-neutral/svg?seed=vibe") 
            )
            
            with gr.Row(equal_height=True):
                msg = gr.Textbox(
                    label="What's your goal?", 
                    placeholder="Type 'Hello' to start, or ask: 'How do lists work?'",
                    lines=1,
                    scale=5,
                    container=False,
                    autofocus=True
                )
                submit_btn = gr.Button("🚀 Start", variant="primary", scale=1)

            # Quick Examples
            gr.Examples(
                examples=[
                    "Hello! I'm new to Python.",
                    "How do for-loops work?",
                    "Explain dictionaries with an example.",
                    "Write a script to calculate Fibonacci numbers."
                ],
                inputs=msg
            )

        # Right Column: Resources Dashboard (Side View - Visible by default)
        with gr.Column(scale=2) as right_col:
            gr.Markdown("### 🎒 Learning Dashboard")
            
            # Tabbed interface for cleaner look
            with gr.Tabs():
                with gr.TabItem("📺 Videos"):
                    video_box_side = gr.Markdown(value="### Recommended Videos\n*Ask a topic to see video suggestions!*")
                
                with gr.TabItem("📚 Reading"):
                    article_box_side = gr.Markdown(value="### Articles & Docs\n*Ask a topic to see reading materials!*")
                
                with gr.TabItem("🧠 Quiz"):
                    quiz_box_side = gr.Markdown(value="### Knowledge Check\n*Ask a topic to unlock the quiz!*")

    # Bottom Row: Resources Dashboard (Bottom View - Hidden by default)
    with gr.Row(visible=False) as bottom_dashboard:
        with gr.Column():
            gr.Markdown("### 🎒 Learning Dashboard")
            with gr.Tabs():
                with gr.TabItem("📺 Videos"):
                    video_box_bottom = gr.Markdown(value="### Recommended Videos\n*Ask a topic to see video suggestions!*")
                
                with gr.TabItem("📚 Reading"):
                    article_box_bottom = gr.Markdown(value="### Articles & Docs\n*Ask a topic to see reading materials!*")
                
                with gr.TabItem("🧠 Quiz"):
                    quiz_box_bottom = gr.Markdown(value="### Knowledge Check\n*Ask a topic to unlock the quiz!*")

    # --- Interaction Logic ---
    async def respond(user_message, history):
        if history is None: history = []
        
        # Immediate user update
        history.append({"role": "user", "content": user_message})
        # Placeholder for AI
        history.append({"role": "assistant", "content": "Thinking..."}) 
        
        # Yield placeholders to ALL output boxes (Side AND Bottom)
        # Structure: history, msg, 3x Side Boxes, 3x Bottom Boxes
        yield history, "", "", "", "", "", "", ""
        
        # Run Agent
        chat_text, video_text, article_text, quiz_text = await run_tutor_dashboard(user_message)
        
        # Update AI response
        history[-1]["content"] = chat_text
        
        # Yield final content to ALL output boxes
        yield history, "", video_text, article_text, quiz_text, video_text, article_text, quiz_text

    # --- Focus Mode Logic ---
    is_fullscreen = gr.State(False)

    def toggle_fullscreen(current_state):
        new_state = not current_state
        # If new_state is True (Fullscreen): Hide side col, Show bottom row
        # If False (Normal): Show side col, Hide bottom row
        
        side_visible = not new_state
        bottom_visible = new_state
        
        btn_text = "↩ Exit Focus" if new_state else "⛶ Focus Mode"
        
        return new_state, gr.Column(visible=side_visible), gr.Row(visible=bottom_visible), btn_text

    fullscreen_btn.click(
        toggle_fullscreen,
        inputs=[is_fullscreen],
        outputs=[is_fullscreen, right_col, bottom_dashboard, fullscreen_btn]
    )

    # Actions - We must map outputs to BOTH side and bottom components
    outputs_list = [
        chatbot, msg, 
        video_box_side, article_box_side, quiz_box_side, 
        video_box_bottom, article_box_bottom, quiz_box_bottom
    ]

    submit_btn.click(
        respond, 
        [msg, chatbot], 
        outputs_list
    )
    
    msg.submit(
        respond, 
        [msg, chatbot], 
        outputs_list
    )

# --- Launch ---
if __name__ == "__main__":
    demo.queue().launch()