import gradio as gr
import subprocess
import threading
import os
import argparse

# Example theme map for the argument parser (add more themes if you have them)
theme_map = {
    "Ocean": gr.themes.Ocean(),
    "Default": gr.themes.Default(),
    # Add other themes as needed
}

# Function to run VNC setup (generates .vnc/passwd)
def setup_vnc(password):
    vnc_dir = os.path.expanduser("~/.vnc")
    os.makedirs(vnc_dir, exist_ok=True)
    passwd_file = os.path.join(vnc_dir, "passwd")
    proc = subprocess.Popen(
        f"echo '{password}' | vncpasswd -f > '{passwd_file}'",
        shell=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    stdout, stderr = proc.communicate()
    os.chmod(passwd_file, 0o600)
    return f"VNC password set to: {password}\n{stdout.decode()}\n{stderr.decode()}"

# Function to run the webui.py application
def start_webui(ip="0.0.0.0", port=7860):
    # This will run in a thread
    def run():
        subprocess.run(["python", "webui.py", "--ip", ip, "--port", str(port)], cwd="/app")
    thread = threading.Thread(target=run, daemon=True)
    thread.start()
    return f"webui.py started at http://{ip}:{port}"

# Gradio interface for VNC setup
def vnc_setup(password):
    result = setup_vnc(password)
    return result

# Main app UI (tab content)
def create_ui(theme_name="Ocean"):
    with gr.Blocks(theme=theme_map.get(theme_name, gr.themes.Ocean())) as base_ui:
        gr.Markdown("## WebUI and VNC Setup")

        with gr.Tab("Start WebUI"):
            with gr.Row():
                ip = gr.Textbox(label="IP", value="0.0.0.0")
                port = gr.Number(label="Port", value=7788)
            webui_btn = gr.Button("Start WebUI")
            webui_out = gr.Textbox(label="WebUI Output")
            webui_btn.click(start_webui, inputs=[ip, port], outputs=webui_out)

        with gr.Tab("VNC Setup"):
            vnc_password = gr.Textbox(label="VNC Password", type="password")
            vnc_btn = gr.Button("Setup VNC Password")
            vnc_out = gr.Textbox(label="VNC Setup Output")
            vnc_btn.click(vnc_setup, inputs=[vnc_password], outputs=vnc_out)
    return base_ui

def main():
    parser = argparse.ArgumentParser(description="Gradio WebUI for Browser Agent")
    parser.add_argument("--ip", type=str, default=os.getenv("HOST", "0.0.0.0"), help="IP address to bind to")
    parser.add_argument("--port", type=int, default=int(os.getenv("PORT", "7860")), help="Port to listen on")
    parser.add_argument("--theme", type=str, default=os.getenv("THEME", "Ocean"), choices=theme_map.keys(), help="Theme to use for the UI")
    args = parser.parse_args()

    # Existing UI
    base_ui = create_ui(theme_name=args.theme)

    # Combine into single Gradio app
    with gr.Blocks(theme=theme_map[args.theme]) as demo:
        gr.Markdown("# 🌐 Browser Agent Web-UI")

        with gr.Tab("Main App"):
            base_ui.render()

        with gr.Tab("Web VNC"):
            # The iframe must point to the actual noVNC server (default: localhost:6080/vnc.html)
            gr.HTML('<iframe src="http://localhost:6080/vnc.html" width="100%" height="800" style="border:none;"></iframe>')

    demo.queue().launch(server_name=args.ip, server_port=args.port)

if __name__ == "__main__":
    main()