import os
import json
import subprocess
from pathlib import Path

import spaces
import gradio as gr


# ============================================================
# ZeroGPU registration
# ============================================================

@spaces.GPU(duration=1)
def _zerogpu_probe():
    return None


# A real Gradio event bound to the GPU-decorated function.
# The button is never exposed to the user.
with gr.Blocks() as _zero_gpu_demo:
    _probe = gr.Button(visible=False)
    _probe.click(
        fn=_zerogpu_probe,
        inputs=[],
        outputs=[],
    )

# IMPORTANT:
# ZeroGPU detects GPU handlers when Blocks.launch() is called.
# We launch Gradio on an INTERNAL port so nanobot can own 7860.
_zero_gpu_demo.launch(
    server_name="127.0.0.1",
    server_port=7861,
    prevent_thread_lock=True,
)


# ============================================================
# Paths
# ============================================================

DATA = Path("/data/nanobot")
CONFIG = DATA / "config.json"
WORKSPACE = DATA / "workspace"

DATA.mkdir(parents=True, exist_ok=True)
WORKSPACE.mkdir(parents=True, exist_ok=True)


# ============================================================
# Secrets
# ============================================================

WEB_TOKEN = os.environ.get("NANOBOT_WEB_TOKEN")
GEMINI_KEY = os.environ.get("GEMINI_API_KEY")

if not WEB_TOKEN:
    raise RuntimeError("NANOBOT_WEB_TOKEN is not set.")

if not GEMINI_KEY:
    raise RuntimeError("GEMINI_API_KEY is not set.")


# ============================================================
# Load existing configuration
# ============================================================

if CONFIG.exists():
    config = json.loads(
        CONFIG.read_text(encoding="utf-8")
    )
else:
    config = {}


# ============================================================
# Gemini provider
# ============================================================

config.setdefault("providers", {})

config["providers"]["gemini"] = {
    "apiKey": "${GEMINI_API_KEY}"
}


# ============================================================
# Primary model
# ============================================================

config.setdefault("modelPresets", {})

config["modelPresets"]["primary"] = {
    "provider": "gemini",
    "model": "gemini-2.5-flash",
    "maxTokens": 8192,
    "contextWindowTokens": 1048576,
    "temperature": 0.1,
}


# ============================================================
# Agent defaults
# ============================================================

config.setdefault("agents", {})
config["agents"].setdefault("defaults", {})

config["agents"]["defaults"]["modelPreset"] = "primary"
config["agents"]["defaults"]["workspace"] = str(WORKSPACE)


# ============================================================
# Nanobot WebUI
# ============================================================

config.setdefault("channels", {})

config["channels"]["websocket"] = {
    "enabled": True,
    "host": "0.0.0.0",
    "port": 7860,
    "tokenIssueSecret": "${NANOBOT_WEB_TOKEN}",
    "websocketRequiresToken": True,
}


# ============================================================
# Tool security
# ============================================================

config.setdefault("tools", {})

config["tools"]["restrictToWorkspace"] = True
config["tools"]["webuiAllowRemotePackageInstall"] = False


# ============================================================
# Save configuration
# ============================================================

CONFIG.write_text(
    json.dumps(config, indent=2),
    encoding="utf-8",
)


# ============================================================
# Start nanobot
# ============================================================

subprocess.run(
    [
        "nanobot",
        "gateway",
        "--config",
        str(CONFIG),
        "--workspace",
        str(WORKSPACE),
    ],
    check=True,
)