#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "textual>=0.75.0",
#     "rich>=13.0.0",
# ]
# ///

from __future__ import annotations

import os
import signal
import subprocess
import urllib.request
import json
from pathlib import Path
from typing import Optional

from rich.text import Text
from textual.app import App, ComposeResult
from textual.containers import Container, Horizontal, Vertical, VerticalScroll
from textual.binding import Binding
from textual.widgets import (
    Button,
    Checkbox,
    Footer,
    Header,
    Input,
    Label,
    Log,
    RadioSet,
    RadioButton,
    Select,
    Static,
    TabbedContent,
    TabPane,
)

SCRIPT_DIR = Path(__file__).resolve().parent.parent
MODELS_DIR = SCRIPT_DIR / "models"
TEMPLATES_DIR = SCRIPT_DIR / "templates"
LOGS_DIR = SCRIPT_DIR / "logs"
RUNNER_BIN = SCRIPT_DIR / "bin" / "llm-runner"

TOTAL_USABLE_VRAM_GB = 30.0

PRESETS = {
    "cluster": {
        "name": "🌟 Multi-Agent Cluster (SuperQwen 27B + 2x Qwen 4B)",
        "desc": "Planner on 11434, Worker 1 on 11435, Worker 2 on 11436. Optimized for Pi & Hermes agents.",
        "port": 11434,
        "ctx": 65536,
        "temp": 0.6,
        "top_p": 0.95,
        "top_k": 40,
        "min_p": 0.05,
        "ngl": 99,
        "flash_attn": True,
        "kv_quant": "q8_0",
        "command": "cluster",
    },
    "muse": {
        "name": "🔮 Muse-Glimmer 30B Solo (128K Context)",
        "desc": "Dedicated deep-reasoning agent with max 128K context on port 11434.",
        "port": 11434,
        "ctx": 131072,
        "temp": 0.6,
        "top_p": 0.95,
        "top_k": 40,
        "min_p": 0.05,
        "ngl": 99,
        "flash_attn": True,
        "kv_quant": "q8_0",
        "command": "muse",
    },
    "superqwen": {
        "name": "🧠 SuperQwen 3.8 27B Solo (256K Context)",
        "desc": "Dedicated ultra-long reasoning agent with max 256K context on port 11434.",
        "port": 11434,
        "ctx": 262144,
        "temp": 0.6,
        "top_p": 0.95,
        "top_k": 40,
        "min_p": 0.05,
        "ngl": 99,
        "flash_attn": True,
        "kv_quant": "q8_0",
        "command": "superqwen",
    },
    "qwen4b": {
        "name": "⚡ Qwen 3.5 4B Solo (Fast Instruct)",
        "desc": "Lightweight high-speed instruct model on port 11434 with ~3.9GB VRAM footprint.",
        "port": 11434,
        "ctx": 32768,
        "temp": 0.7,
        "top_p": 0.8,
        "top_k": 20,
        "min_p": 0.05,
        "ngl": 99,
        "flash_attn": True,
        "kv_quant": "q8_0",
        "command": "qwen4b",
    },
    "custom": {
        "name": "🛠️ Custom Model & Parameters",
        "desc": "Specify exact GGUF weight files, context window, and sampling hyperparameters.",
        "port": 11434,
        "ctx": 65536,
        "temp": 0.6,
        "top_p": 0.95,
        "top_k": 40,
        "min_p": 0.05,
        "ngl": 99,
        "flash_attn": True,
        "kv_quant": "q8_0",
        "command": "custom",
    },
}


def find_gguf_models() -> list[tuple[str, str]]:
    models: list[tuple[str, str]] = []
    if MODELS_DIR.exists():
        for path in sorted(MODELS_DIR.rglob("*.gguf")):
            if "mmproj" not in path.name.lower():
                rel = path.relative_to(MODELS_DIR)
                models.append((f"{path.name} ({path.stat().st_size / 1e9:.1f} GB)", str(path)))
    if not models:
        models.append(("No GGUF models found in models/", ""))
    return models


def find_mmproj_files() -> list[tuple[str, str]]:
    projectors = [("None (Text Only)", "")]
    if MODELS_DIR.exists():
        for path in sorted(MODELS_DIR.rglob("mmproj*.gguf")):
            projectors.append((f"{path.name} ({path.stat().st_size / 1e9:.1f} GB)", str(path)))
    return projectors


def check_port_status(port: int) -> tuple[bool, str]:
    url = f"http://127.0.0.1:{port}/health"
    try:
        req = urllib.request.Request(url, headers={"User-Agent": "LLM-TUI"})
        with urllib.request.urlopen(req, timeout=0.8) as response:
            data = json.loads(response.read().decode())
            status = data.get("status", "unknown")
            return True, f"Online ({status})"
    except Exception:
        return False, "Offline"


class LLMConfigApp(App):
    CSS = """
    Screen {
        background: #14171d;
        color: #d8dee9;
    }
    Header {
        background: #1e222b;
        color: #88c0d0;
        text-style: bold;
    }
    TabbedContent {
        height: 1fr;
    }
    TabPane {
        padding: 1 2;
    }
    .card {
        background: #1a1e27;
        border: solid #3b4252;
        padding: 1 2;
        margin-bottom: 1;
    }
    .title {
        color: #88c0d0;
        text-style: bold;
        margin-bottom: 1;
    }
    .field-row {
        height: auto;
        margin-bottom: 1;
    }
    .field-label {
        width: 18;
        color: #81a1c1;
        padding-top: 1;
    }
    Input {
        background: #232834;
        border: solid #434c5e;
        color: #eceff4;
        width: 1fr;
    }
    Input:focus {
        border: solid #88c0d0;
    }
    Select {
        background: #232834;
        border: solid #434c5e;
        color: #eceff4;
        width: 1fr;
    }
    Button {
        margin-right: 1;
        background: #2e3440;
        color: #eceff4;
        border: solid #4c566a;
    }
    Button:hover {
        background: #434c5e;
    }
    Button.-primary {
        background: #5e81ac;
        color: #eceff4;
        text-style: bold;
        border: solid #81a1c1;
    }
    Button.-error {
        background: #bf616a;
        color: #eceff4;
        text-style: bold;
        border: solid #d08770;
    }
    #vram-meter {
        background: #232834;
        border: solid #4c566a;
        padding: 1;
        margin-top: 1;
    }
    #status-card {
        background: #1e222b;
        border: solid #4c566a;
        padding: 1;
        margin-bottom: 1;
    }
    Log {
        background: #0f1117;
        border: solid #3b4252;
        height: 1fr;
    }
    """

    BINDINGS = [
        Binding("f5", "launch", "Launch Server"),
        Binding("f8", "stop", "Stop Selected"),
        Binding("f9", "stop_all", "Stop All"),
        Binding("f2", "refresh_status", "Refresh Status"),
        Binding("ctrl+c", "quit", "Quit"),
        Binding("q", "quit", "Quit"),
    ]

    selected_preset: str = "cluster"
    models_list: list[tuple[str, str]] = []
    mmproj_list: list[tuple[str, str]] = []

    def compose(self) -> ComposeResult:
        yield Header(show_clock=True)
        self.models_list = find_gguf_models()
        self.mmproj_list = find_mmproj_files()

        with TabbedContent():
            with TabPane("🚀 Model Launcher & Form", id="tab-launch"):
                with VerticalScroll():
                    # 1. Preset Selector
                    with Container(classes="card"):
                        yield Label("Select Inference Preset", classes="title")
                        with RadioSet(id="preset-radios"):
                            for key, preset in PRESETS.items():
                                yield RadioButton(
                                    preset["name"],
                                    id=f"preset-{key}",
                                    value=(key == "cluster"),
                                )
                        yield Label(
                            PRESETS["cluster"]["desc"],
                            id="preset-description",
                            classes="text-muted",
                        )

                    # 2. Parameters Form
                    with Container(classes="card", id="form-container"):
                        yield Label("Hyperparameters & Hardware Tuning", classes="title")

                        with Horizontal(classes="field-row"):
                            yield Label("GGUF Model:", classes="field-label")
                            yield Select(
                                options=[(lbl, val) for lbl, val in self.models_list],
                                value=self.models_list[0][1] if self.models_list else None,
                                id="select-model",
                            )

                        with Horizontal(classes="field-row"):
                            yield Label("Vision MMProj:", classes="field-label")
                            yield Select(
                                options=[(lbl, val) for lbl, val in self.mmproj_list],
                                value=self.mmproj_list[0][1] if self.mmproj_list else None,
                                id="select-mmproj",
                            )

                        with Horizontal(classes="field-row"):
                            yield Label("Context Size:", classes="field-label")
                            yield Input(value="65536", id="input-ctx")

                        with Horizontal(classes="field-row"):
                            yield Label("Port:", classes="field-label")
                            yield Input(value="11434", id="input-port")

                        with Horizontal(classes="field-row"):
                            yield Label("Temperature:", classes="field-label")
                            yield Input(value="0.6", id="input-temp")

                        with Horizontal(classes="field-row"):
                            yield Label("Top-P:", classes="field-label")
                            yield Input(value="0.95", id="input-top-p")

                        with Horizontal(classes="field-row"):
                            yield Label("Min-P:", classes="field-label")
                            yield Input(value="0.05", id="input-min-p")

                        with Horizontal(classes="field-row"):
                            yield Label("GPU Layers (-ngl):", classes="field-label")
                            yield Input(value="99", id="input-ngl")

                        with Horizontal(classes="field-row"):
                            yield Label("KV Cache Quant:", classes="field-label")
                            yield Select(
                                options=[("q8_0 (High Quality, 50% savings)", "q8_0"), ("f16 (Full Precision)", "f16"), ("q4_0 (Max Memory Savings)", "q4_0")],
                                value="q8_0",
                                id="select-kv",
                            )

                        with Horizontal(classes="field-row"):
                            yield Label("Flash Attention:", classes="field-label")
                            yield Checkbox("Enabled (-fa on)", value=True, id="check-fa")

                        # VRAM Estimation Banner
                        yield Static(id="vram-meter")

                    # 3. Execution Action Bar
                    with Horizontal(classes="card"):
                        yield Button("🚀 Launch Server (F5)", variant="primary", id="btn-launch")
                        yield Button("🛑 Stop Server (F8)", variant="error", id="btn-stop")
                        yield Button("⚠️ Stop All Instances (F9)", variant="error", id="btn-stop-all")
                        yield Button("🔄 Refresh Health (F2)", id="btn-refresh")

            with TabPane("📊 Server Status & Live Logs", id="tab-status"):
                with Vertical():
                    yield Static(id="status-card")
                    yield Label("Live Logs Streaming:", classes="title")
                    yield Log(id="log-stream", highlight=True)

        yield Footer()

    def on_mount(self) -> None:
        self.update_vram_estimation()
        self.action_refresh_status()
        self.set_interval(3.0, self.action_refresh_status)
        self.set_interval(1.5, self.tail_logs)

    def on_radio_set_changed(self, event: RadioSet.Changed) -> None:
        radio_id = event.pressed.id or ""
        preset_key = radio_id.replace("preset-", "")
        if preset_key in PRESETS:
            self.selected_preset = preset_key
            preset = PRESETS[preset_key]
            self.query_one("#preset-description", Label).update(preset["desc"])
            self.query_one("#input-port", Input).value = str(preset["port"])
            self.query_one("#input-ctx", Input).value = str(preset["ctx"])
            self.query_one("#input-temp", Input).value = str(preset["temp"])
            self.query_one("#input-top-p", Input).value = str(preset["top_p"])
            self.query_one("#input-min-p", Input).value = str(preset["min_p"])
            self.query_one("#input-ngl", Input).value = str(preset["ngl"])
            self.query_one("#check-fa", Checkbox).value = preset["flash_attn"]
            self.query_one("#select-kv", Select).value = preset["kv_quant"]
            self.update_vram_estimation()

    def on_input_changed(self, event: Input.Changed) -> None:
        self.update_vram_estimation()

    def on_select_changed(self, event: Select.Changed) -> None:
        self.update_vram_estimation()

    def update_vram_estimation(self) -> None:
        try:
            ctx = int(self.query_one("#input-ctx", Input).value or "0")
        except ValueError:
            ctx = 65536

        preset = self.selected_preset
        if preset == "cluster":
            weights_gb = 25.9
            kv_gb = 3.2
        elif preset == "muse":
            weights_gb = 17.5
            kv_gb = (ctx / 131072.0) * 8.5
        elif preset == "superqwen":
            weights_gb = 17.5
            kv_gb = (ctx / 262144.0) * 10.5
        elif preset == "qwen4b":
            weights_gb = 3.9
            kv_gb = (ctx / 32768.0) * 1.5
        else:
            model_val = str(self.query_one("#select-model", Select).value or "")
            if model_val and Path(model_val).exists():
                weights_gb = Path(model_val).stat().st_size / 1e9
            else:
                weights_gb = 15.0
            kv_gb = (ctx / 65536.0) * 4.0

        total_est = weights_gb + kv_gb
        remaining = TOTAL_USABLE_VRAM_GB - total_est

        if total_est <= TOTAL_USABLE_VRAM_GB:
            status_color = "green"
            verdict = "SAFE (Fits in VRAM)"
        else:
            status_color = "red"
            verdict = "OVERFLOW WARNING (May spill to system RAM / ROCm crash)"

        meter_text = (
            f"VRAM Budget (~{TOTAL_USABLE_VRAM_GB:.0f}GB Usable): "
            f"Weights ~{weights_gb:.1f}GB + KV Cache ~{kv_gb:.1f}GB = "
            f"[{status_color}]{total_est:.1f}GB total[/] | Remaining Headroom: ~{remaining:.1f}GB [{verdict}]"
        )
        self.query_one("#vram-meter", Static).update(meter_text)

    def action_refresh_status(self) -> None:
        lines = []
        for port, label in [(11434, "Planner / Solo"), (11435, "Worker 1"), (11436, "Worker 2")]:
            online, msg = check_port_status(port)
            badge = "[green]● ONLINE[/]" if online else "[yellow]○ OFFLINE[/]"
            lines.append(f"Port {port} ({label}): {badge} - {msg}")

        card = self.query_one("#status-card", Static)
        card.update("\n".join(lines))

    def tail_logs(self) -> None:
        log_widget = self.query_one("#log-stream", Log)
        logs = sorted(LOGS_DIR.glob("*.log"), key=os.path.getmtime, reverse=True)
        if logs:
            latest = logs[0]
            try:
                with open(latest, "r", errors="ignore") as f:
                    f.seek(max(0, latest.stat().st_size - 8192))
                    tail = f.read()
                    log_widget.clear()
                    log_widget.write(f"--- Streaming {latest.name} ---\n" + tail)
            except Exception:
                pass

    def action_launch(self) -> None:
        preset = self.selected_preset
        log_widget = self.query_one("#log-stream", Log)
        log_widget.write(f"\n[bold cyan]>>> Launching preset: {preset}...[/]\n")

        env = os.environ.copy()
        env["CTX"] = self.query_one("#input-ctx", Input).value or "65536"
        env["PORT"] = self.query_one("#input-port", Input).value or "11434"
        env["TEMP"] = self.query_one("#input-temp", Input).value or "0.6"
        env["TOP_P"] = self.query_one("#input-top-p", Input).value or "0.95"
        env["MIN_P"] = self.query_one("#input-min-p", Input).value or "0.05"
        env["NGL"] = self.query_one("#input-ngl", Input).value or "99"
        env["FLASH_ATTN"] = "on" if self.query_one("#check-fa", Checkbox).value else "off"
        env["CTK"] = str(self.query_one("#select-kv", Select).value or "q8_0")
        env["CTV"] = env["CTK"]

        cmd = [str(RUNNER_BIN), preset]
        subprocess.Popen(cmd, env=env, cwd=str(SCRIPT_DIR), start_new_session=True)
        self.notify(f"Launching {preset} server in background. Checking ports...", severity="information")
        self.action_refresh_status()

    def action_stop(self) -> None:
        port = self.query_one("#input-port", Input).value or "11434"
        subprocess.run([str(RUNNER_BIN), "stop", port], cwd=str(SCRIPT_DIR))
        self.notify(f"Stopped server on port {port}", severity="warning")
        self.action_refresh_status()

    def action_stop_all(self) -> None:
        subprocess.run([str(RUNNER_BIN), "stop", "all"], cwd=str(SCRIPT_DIR))
        self.notify("Stopped all running LLM servers", severity="warning")
        self.action_refresh_status()

    def on_button_pressed(self, event: Button.Pressed) -> None:
        if event.button.id == "btn-launch":
            self.action_launch()
        elif event.button.id == "btn-stop":
            self.action_stop()
        elif event.button.id == "btn-stop-all":
            self.action_stop_all()
        elif event.button.id == "btn-refresh":
            self.action_refresh_status()


if __name__ == "__main__":
    app = LLMConfigApp()
    app.run()
