import gradio as gr
import os
import logging
import re
import sys
import io
import json
import time
import base64
import tempfile
import subprocess
from datetime import datetime, timedelta
from openai import OpenAI

# Try importing E2B
try:
    from e2b_code_interpreter import Sandbox
    E2B_AVAILABLE = True
except ImportError:
    E2B_AVAILABLE = False

# Import for Web Search
try:
    from ddgs import DDGS
    SEARCH_AVAILABLE = True
except ImportError:
    SEARCH_AVAILABLE = False

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("FyodorIDE")

# ==================== CONFIGURATION ====================

MINIMAX_API_KEY = os.getenv("MINIMAX_API_KEY")
E2B_API_KEY = os.getenv("E2B_API_KEY")

API_CONFIG = {
    "api_key": MINIMAX_API_KEY if MINIMAX_API_KEY else "dummy-key",
    "base_url": "https://api.minimax.io/v1",
}
MODEL_ID = "MiniMax-M2"
SESSION_TIMEOUT_SEC = 15 * 60

# ==================== HELPER FUNCTIONS ====================

def clean_thought_process(text):
    """Removes the <think>...</think> blocks from the text to prevent leaks."""
    if not text: return ""
    # Remove completed blocks
    text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
    # Remove pending/unclosed blocks
    if "<think>" in text:
        text = text.split("<think>")[0]
    return text.strip()

def export_script(code):
    if not code.strip(): return None
    try:
        with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.py', prefix='fyodor_script_') as f:
            f.write(code)
            return f.name
    except Exception as e:
        return None

def detect_language(filename):
    """Returns the language for the code editor based on extension."""
    if filename.endswith(".py"): return "python"
    if filename.endswith(".cpp") or filename.endswith(".h") or filename.endswith(".hpp"): return "cpp"
    if filename.endswith(".js"): return "javascript"
    if filename.endswith(".html"): return "html"
    if filename.endswith(".css"): return "css"
    return "python"

def robust_search(query):
    if not SEARCH_AVAILABLE: return "Error: Missing duckduckgo-search"
    try:
        results = []
        with DDGS() as ddgs:
            ddg_gen = ddgs.text(query, max_results=5)
            if ddg_gen:
                for r in ddg_gen:
                    results.append(f"### [{r.get('title')}]({r.get('href')})\n{r.get('body')}\n")
        
        if not results: return "⚠️ No results found."
        return "\n".join(results)
    except Exception as e:
        return f"❌ Search Error: {str(e)}"

# ==================== TOOL: E2B SANDBOX (POLYGLOT) ====================

def execute_code_e2b(project_state, mode="Code Mode"):
    """
    Polymorphic execution engine with Mode-Gated Internet Access.
    """
    if not E2B_AVAILABLE or not E2B_API_KEY:
        return "❌ Error: E2B API Key missing. Cannot execute code."

    enable_internet = (mode == "Web Mode")
    logger.info(f"⚡ Executing in E2B Sandbox (Mode: {mode}, Internet: {enable_internet})...")
    
    try:
        with Sandbox.create(allow_internet_access=enable_internet) as sandbox:
            # 1. Write all auxiliary files to the sandbox filesystem
            for filename, content in project_state.items():
                if filename not in ["main.py", "main.cpp"]: 
                    sandbox.files.write(filename, content)
            
            output_log = ""
            net_status = "🟢 ONLINE" if enable_internet else "🔒 OFFLINE"
            output_log += f"**Environment Status:** {net_status}\n\n"

            # 2. C++ EXECUTION PATH (Wrapper Method)
            if "main.cpp" in project_state:
                logger.info("⚙️ Detected C++ Project. Compiling via Python wrapper...")
                
                # Write C++ source explicitly ensures it's fresh
                sandbox.files.write("main.cpp", project_state["main.cpp"])
                
                # Trojan Horse: Use Python to manage the compilation process reliably
                cpp_runner_script = """
import subprocess
import sys

print("⚙️ **Compiling C++...**")
compile_res = subprocess.run(["g++", "-o", "main", "main.cpp"], capture_output=True, text=True)

if compile_res.returncode != 0:
    print("❌ **COMPILATION FAILED**")
    print(compile_res.stderr)
else:
    print("🚀 **Running Binary...**")
    run_res = subprocess.run(["./main"], capture_output=True, text=True)
    print("📄 **STDOUT**:")
    print(run_res.stdout)
    if run_res.stderr:
        print("⚠️ **STDERR**:")
        print(run_res.stderr)
"""
                execution = sandbox.run_code(cpp_runner_script)
                
                # Properly join list of strings into a single string
                if execution.logs.stdout:
                    output_log += "".join(execution.logs.stdout) + "\n"
                if execution.logs.stderr:
                    output_log += f"⚠️ **INTERNAL ERROR**:\n{''.join(execution.logs.stderr)}\n"
                
                return output_log.strip()

            # 3. PYTHON EXECUTION PATH (Default)
            elif "main.py" in project_state:
                code_to_run = project_state["main.py"]
                sandbox.files.write("main.py", code_to_run)
                
                execution = sandbox.run_code(code_to_run)
                
                if execution.logs.stdout:
                    log_str = "".join(execution.logs.stdout)
                    output_log += f"📄 **STDOUT**:\n```\n{log_str}\n```"
                if execution.logs.stderr:
                    err_str = "".join(execution.logs.stderr)
                    output_log += f"\n⚠️ **STDERR**:\n```\n{err_str}\n```"
                
                for result in execution.results:
                    if hasattr(result, 'png') and result.png:
                        img_md = f"![Plot](data:image/png;base64,{result.png})"
                        output_log += f"\n🖼️ **Plot**:\n{img_md}"
                
                return output_log if output_log.strip() else f"✅ Executed (No Output)\n\n{output_log}"
            
            else:
                return "⚠️ No executable entry point found (main.py or main.cpp)."

    except Exception as e:
        return f"❌ Sandbox Error: {str(e)}"

# ==================== CORE LOGIC ====================

def init_session():
    return {"history": [], "start_time": time.time(), "active": True}

def get_system_prompt(mode):
    base_prompt = """
You are Fyodor, an advanced AI-Native IDE engine.
You are running in a FULLY ISOLATED E2B Sandbox (No Internet).
You have access to a Multi-File Project environment.

FILESYSTEM RULES:
- The entry point is 'main.py' (Python) OR 'main.cpp' (C++).
- You can create helper modules/headers (e.g., 'utils.py', 'math.h').
- When providing code, specify the filename if strictly necessary.

CRITICAL INSTRUCTIONS:
- CODE: You MUST use standard Markdown code blocks: ```python ... ``` or ```cpp ... ```.
- DO NOT use XML tags. DO NOT use JSON for code delivery.
- JUST WRITE THE CODE BLOCK.
"""
    if mode == "Data Mode":
        return base_prompt + "\n\nMODE: DATA ANALYSIS\n- Focus on pandas, numpy, and statistical insights.\n- NO INTERNET ACCESS."
    elif mode == "Visualization Mode":
        return base_prompt + "\n\nMODE: VISUALIZATION\n- Focus on matplotlib/seaborn.\n- Create visually striking plots.\n- NO INTERNET ACCESS."
    elif mode == "Web Mode":
        return base_prompt + "\n\nMODE: WEB ARCHITECT\n- Focus on modern web technologies (Flask, HTML, CSS).\n- **YOU HAVE INTERNET ACCESS**. Use `requests` to fetch real data.\n- Structure code for web servers or API clients."
    elif mode == "C++ Mode":
        return base_prompt + "\n\nMODE: C++ DEVELOPER\n- Focus on modern C++ (17/20).\n- Use iostream, vector, string.\n- Ensure code compiles with g++.\n- NO INTERNET ACCESS."
    else: # Code Mode
        return base_prompt + "\n\nMODE: GENERAL CODING\n- Focus on clean, modular Python/C++ code.\n- NO INTERNET ACCESS."

def parse_code_update(text, current_project, active_file):
    """
    Parses LLM output for code blocks. 
    Updates the ACTIVE file with the last code block found.
    """
    clean_text = clean_thought_process(text)
    
    # Robust Regex for code blocks (Python, C++, JS, HTML, CSS)
    # Using multi-line string construction to prevent truncation errors
    pattern = (
        r"```"          # Opening backticks
        r"(?:python|cpp|c\+\+|javascript|html|css)?"  # Optional language identifier
        r"\s*\n"        # Optional whitespace and mandatory newline
        r"(.*?)"        # The code content (non-greedy)
        r"```"          # Closing backticks
    )
    
    matches = re.findall(pattern, clean_text, re.DOTALL)
    
    if matches:
        new_code = matches[-1]
        current_project[active_file] = new_code
        return new_code, current_project
        
    return None, current_project

def process_query(user_input, session_data, project_state, active_file, mode):
    """Main Chat Loop."""
    if session_data is None: session_data = init_session()
    
    # Auto-renewal check
    if time.time() - session_data["start_time"] > SESSION_TIMEOUT_SEC:
        session_data = init_session()

    history = session_data["history"] 
    
    if not user_input:
        yield history, project_state, project_state[active_file], session_data
        return

    # Update Internal State
    history.append({"role": "user", "content": user_input})
    history.append({"role": "assistant", "content": f"🧠 *Fyodor is thinking in {mode}...*"})
    
    # Yield UI Update
    yield history, project_state, project_state[active_file], session_data

    # Context Construction
    sys_prompt = get_system_prompt(mode)
    
    project_context = "\n\n=== CURRENT PROJECT FILES ===\n"
    for fname, content in project_state.items():
        project_context += f"--- {fname} ---\n{content}\n"
    
    messages = [{"role": "system", "content": sys_prompt + project_context}]

    # Build Chat History for API
    for msg in history[:-2][-8:]:
        if isinstance(msg, dict):
            if msg['role'] == 'user':
                messages.append({"role": "user", "content": msg['content']})
            elif msg['role'] == 'assistant':
                clean_content = msg['content'].split("---")[0].strip()
                messages.append({"role": "assistant", "content": clean_content})
    
    messages.append({"role": "user", "content": user_input})

    try:
        client = OpenAI(**API_CONFIG)
        response = client.chat.completions.create(
            model=MODEL_ID, messages=messages, temperature=0.1, max_tokens=4000, stream=True
        )

        full_response = ""
        for chunk in response:
            if chunk.choices and chunk.choices[0].delta:
                delta_content = chunk.choices[0].delta.content
                if delta_content:
                    full_response += delta_content
                    clean_display = clean_thought_process(full_response)
                    if isinstance(history[-1], dict):
                        history[-1]['content'] = clean_display
                    yield history, project_state, project_state[active_file], session_data

        new_code, updated_project = parse_code_update(full_response, project_state, active_file)
        
        if new_code:
            if isinstance(history[-1], dict): history[-1]['content'] += f"\n\n📝 *Updated {active_file}*"
            yield history, updated_project, updated_project[active_file], session_data

            if isinstance(history[-1], dict): history[-1]['content'] += "\n\n⏳ *Running Project...*"
            yield history, updated_project, updated_project[active_file], session_data
            
            # PASS MODE TO EXECUTOR
            exec_result = execute_code_e2b(updated_project, mode=mode)
            if isinstance(history[-1], dict): history[-1]['content'] += f"\n\n---\n{exec_result}"
            
            yield history, updated_project, updated_project[active_file], session_data

    except Exception as e:
        if isinstance(history[-1], dict): history[-1]['content'] += f"\n\n❌ Error: {str(e)}"
        yield history, project_state, project_state[active_file], session_data

# ==================== MANUAL ACTION HANDLERS ====================

def manual_run_project(project_state, session_data, mode):
    if session_data is None: session_data = init_session()
    history = session_data["history"]
    history.append({"role": "user", "content": "▶️ **Run Project (Manual)**"})
    result = execute_code_e2b(project_state, mode=mode)
    history.append({"role": "assistant", "content": f"**Execution Result:**\n\n{result}"})
    return history, session_data

def run_code_review(code, session_data):
    if not code.strip(): return session_data["history"], session_data
    history = session_data["history"]
    history.append({"role": "user", "content": "🕵️ **Requesting Code Review**"})
    history.append({"role": "assistant", "content": "🤔 *Reviewing...*"})
    
    messages = [
        {"role": "system", "content": "Senior Architect Code Review. Check for security and efficiency."},
        {"role": "user", "content": code}
    ]
    try:
        client = OpenAI(**API_CONFIG)
        response = client.chat.completions.create(model=MODEL_ID, messages=messages, max_tokens=2000)
        review = response.choices[0].message.content if response.choices else "No review."
        history[-1]["content"] = f"### 🕵️ Code Review Report\n\n{clean_thought_process(review)}"
    except Exception as e:
        history[-1]["content"] = f"❌ Review Failed: {e}"
    return history, session_data

def clear_session_chat(session_data):
    if session_data is None: session_data = init_session()
    session_data["history"] = []
    return [], session_data

def update_active_file(selected_file, project_data):
    lang = detect_language(selected_file)
    return gr.update(value=project_data.get(selected_file, ""), language=lang)

def save_file_change(new_content, selected_file, project_data):
    project_data[selected_file] = new_content
    return project_data

def manual_save_file(new_content, selected_file, project_data):
    project_data[selected_file] = new_content
    return project_data, gr.Info(f"✅ Saved {selected_file} to memory")

def create_new_file(new_name, project_data):
    if not new_name: return gr.update(choices=list(project_data.keys())), project_data, gr.update()
    if new_name in project_data: return gr.update(choices=list(project_data.keys())), project_data, gr.update()
    
    content = ""
    if new_name.endswith(".cpp"):
        content = "#include <iostream>\n\nint main() {\n    std::cout << \"Hello from C++!\" << std::endl;\n    return 0;\n}"
    else:
        content = "# New file"
        
    project_data[new_name] = content
    lang = detect_language(new_name)
    return gr.update(choices=list(project_data.keys()), value=new_name), project_data, gr.update(value=content, language=lang)

def toggle_sidebar(is_visible):
    return not is_visible, gr.update(visible=not is_visible)

# ==================== UI ====================

VSCODE_TERMINAL_CSS = """
:root { --bg-dark: #0C0C0C; --panel-bg: #1e1e1e; --sidebar-bg: #252526; --border-col: #333333; --accent: #3A96DD; --text-main: #CCCCCC; --success: #13A10E; }
body, .gradio-container { background-color: var(--bg-dark) !important; color: var(--text-main) !important; font-family: 'Consolas', monospace !important; padding: 0 !important; }
.main-panel { background-color: var(--panel-bg) !important; border: 1px solid var(--border-col) !important; border-radius: 0px !important; }
.gradio-row { gap: 0 !important; }
.gradio-column { background-color: var(--bg-dark); border-right: 1px solid var(--border-col); padding: 10px !important; }
.chatbot { background-color: var(--bg-dark) !important; border: none !important; height: 100% !important; font-family: 'Consolas', monospace !important; }
.chatbot .message { padding: 8px 12px !important; border-bottom: 1px solid #1a1a1a !important; border-radius: 0 !important; box-shadow: none !important; }
.chatbot .user.message { background-color: #1a1a1a !important; border-left: 2px solid var(--accent) !important; color: #fff !important; }
.chatbot .bot.message { background-color: transparent !important; border-left: 2px solid #333 !important; }
textarea, input[type="text"] { background-color: #3c3c3c !important; border: 1px solid #3c3c3c !important; color: #fff !important; border-radius: 2px !important; }
textarea:focus, input[type="text"]:focus { border-color: var(--accent) !important; box-shadow: 0 0 0 1px var(--accent) !important; }
button { border-radius: 2px !important; text-transform: uppercase; font-weight: 600; letter-spacing: 0.5px; font-size: 12px !important; transition: all 0.2s ease-in-out !important; }
button.primary { background: var(--accent) !important; color: #fff !important; border: none !important; }
button.primary:hover { background-color: #4DA2E5 !important; box-shadow: 0 0 8px rgba(58, 150, 221, 0.4); }
button.secondary { background: #333 !important; color: #fff !important; border: 1px solid #444 !important; }
button.secondary:hover { background-color: #444 !important; border-color: #555 !important; }
.cm-editor { border: 1px solid var(--border-col) !important; border-radius: 0 !important; }
#top-bar { background-color: var(--sidebar-bg); border-bottom: 1px solid var(--border-col); min-height: 50px; padding: 0 10px !important; display: flex; align-items: center; justify-content: space-between; }
h1 { font-size: 1.2rem !important; color: var(--text-main) !important; font-weight: normal !important; margin: 0 !important; padding: 0 !important; background-color: transparent !important; border: none !important; display: flex; align-items: center; }
h1::before { content: "⚡"; margin-right: 8px; color: var(--accent); }
.status-label { background-color: transparent !important; border: none !important; box-shadow: none !important; text-align: right !important; font-family: 'Consolas', monospace !important; font-weight: bold; color: var(--success) !important; width: auto !important; min-width: 200px; }
.status-label textarea, .status-label input { text-align: right !important; padding-right: 10px !important; background-color: transparent !important; border: none !important; box-shadow: none !important; }
.status-bar { background-color: var(--accent) !important; color: white !important; font-size: 12px !important; padding: 2px 10px !important; height: 24px !important; display: flex; align-items: center; border-top: 1px solid var(--border-col); }
.mini-chat { border-top: 1px solid var(--border-col) !important; }
.sidebar-tabs button { background-color: transparent !important; border-bottom: none !important; }
.sidebar-tabs button.selected { color: var(--accent) !important; border-bottom: 2px solid var(--accent) !important; }
footer { display: none !important; }
"""

def run_app():
    with gr.Blocks(title="Fyodor IDE | Pro") as demo:
        session_state = gr.State(init_session())
        project_state = gr.State({"main.py": "import numpy as np\nprint('Hello Fyodor Pro!')"})
        active_file_state = gr.State("main.py")
        sidebar_visible = gr.State(True)

        with gr.Column(elem_classes="main-panel"):
            # Top Bar
            with gr.Row(elem_id="top-bar"):
                with gr.Row(scale=1): gr.Markdown("# FYODOR IDE // PRO")
                with gr.Row(scale=1):
                    sidebar_toggle_btn = gr.Button("👁️ Toggle Sidebar", size="sm", variant="secondary")
                    timer_display = gr.Textbox(value="READY", show_label=False, interactive=False, container=False, elem_classes="status-label")

            # Workspace
            with gr.Row():
                # Editor (Left)
                with gr.Column(scale=4):
                    with gr.Row(elem_classes="file-tabs"):
                        file_dropdown = gr.Dropdown(choices=["main.py"], value="main.py", show_label=False, container=False, scale=2)
                        new_file_txt = gr.Textbox(placeholder="filename.py / filename.cpp", show_label=False, container=False, scale=1)
                        new_file_btn = gr.Button("+", scale=0, min_width=40, variant="secondary")
                    
                    code_editor = gr.Code(language="python", value="import numpy as np\nprint('Hello Fyodor Pro!')", lines=32, interactive=True, show_label=False)
                    
                    with gr.Row():
                        save_btn = gr.Button("SAVE FILE", size="sm", variant="secondary") 
                        run_manual_btn = gr.Button("▶ EXECUTE SCRIPT", variant="primary")
                        review_btn = gr.Button("🕵 REVIEW CODE", size="sm", variant="secondary")

                # Terminal (Right)
                with gr.Column(scale=2):
                    mode_radio = gr.Radio(choices=["Code Mode", "C++ Mode", "Data Mode", "Web Mode"], value="Code Mode", label="SYSTEM MODE", elem_classes="mode-radio")
                    chatbot = gr.Chatbot(label="TERMINAL", height=600, render_markdown=True, elem_classes="chatbot", show_label=False)
                    with gr.Row():
                        msg_input = gr.Textbox(placeholder="Type command...", show_label=False, scale=4, container=False)
                        run_btn = gr.Button("RUN", variant="primary", scale=1)
                        clear_btn = gr.Button("🗑️", variant="secondary", scale=0, min_width=40)

                # RIGHT: SIDEBAR
                with gr.Column(scale=1, min_width=250, visible=True) as sidebar_col:
                    gr.Markdown("### 📚 TOOLS")
                    with gr.Tabs(elem_classes="sidebar-tabs"):
                        with gr.TabItem("🌐 WEB"):
                            mini_chat = gr.Chatbot(label="SEARCH", height=500, elem_classes="mini-chat", show_label=False)
                            with gr.Row():
                                mini_input = gr.Textbox(placeholder="Query...", show_label=False, container=False, scale=3)
                                mini_btn = gr.Button("GO", scale=1, variant="secondary")
                        
                        with gr.TabItem("ℹ️ INFO"):
                            gr.Markdown("""
                            ### 🛡️ Fyodor IDE
                            **Multi-Language Environment**
                            
                            **Languages:**
                            - **Python:** `main.py` (Default)
                            - **C++:** `main.cpp` (Compiles with g++)
                            
                            **Privacy:**
                            - Sandbox: Fully Isolated.
                            - Search: Anonymous.
                            """)

            with gr.Row(elem_classes="status-bar"):
                gr.Markdown("E2B: ONLINE | DUCKDUCKGO: READY | POLYGLOT ENV", elem_id="status-text")

        # Logic
        msg_input.submit(process_query, [msg_input, session_state, project_state, active_file_state, mode_radio], [chatbot, project_state, code_editor, session_state])
        run_btn.click(process_query, [msg_input, session_state, project_state, active_file_state, mode_radio], [chatbot, project_state, code_editor, session_state])

        file_dropdown.change(update_active_file, [file_dropdown, project_state], [code_editor]).then(lambda f: f, [file_dropdown], [active_file_state])
        code_editor.change(save_file_change, [code_editor, active_file_state, project_state], [project_state])
        save_btn.click(manual_save_file, [code_editor, active_file_state, project_state], [project_state])
        new_file_btn.click(create_new_file, [new_file_txt, project_state], [file_dropdown, project_state, code_editor])

        run_manual_btn.click(manual_run_project, [project_state, session_state, mode_radio], [chatbot, session_state])
        review_btn.click(run_code_review, [code_editor, session_state], [chatbot, session_state])
        
        clear_btn.click(clear_session_chat, [session_state], [chatbot, session_state])
        
        # Mini Agent: DICT FORMAT (Gradio 6)
        mini_input.submit(lambda m, h: (h or []) + [{"role": "user", "content": m}, {"role": "assistant", "content": robust_search(m)}], [mini_input, mini_chat], [mini_chat])
        mini_btn.click(lambda m, h: (h or []) + [{"role": "user", "content": m}, {"role": "assistant", "content": robust_search(m)}], [mini_input, mini_chat], [mini_chat])
        
        sidebar_toggle_btn.click(toggle_sidebar, [sidebar_visible], [sidebar_visible, sidebar_col])
        
        def update_timer(session_data):
            if session_data is None: session_data = init_session()
            elapsed = time.time() - session_data["start_time"]
            remaining = max(0, SESSION_TIMEOUT_SEC - elapsed)
            if remaining == 0: session_data = init_session()
            mins, secs = divmod(int(remaining), 60)
            return f"SESSION: {mins}:{secs:02d}", session_data

        timer = gr.Timer(1)
        timer.tick(update_timer, inputs=[session_state], outputs=[timer_display, session_state])

    return demo

if __name__ == "__main__":
    run_app().queue().launch(server_name="0.0.0.0", server_port=7860, css=VSCODE_TERMINAL_CSS)