"""
xorfice CLI — Premium Dynamic TUI Server Setup
==============================================
SOTA Terminal UX with arrow-key navigation and rich live updates.
"""

import argparse
import json
import os
import subprocess
import sys
import threading
import time
import urllib.error
import urllib.request

from rich.console import Console
from rich.live import Live
from rich.panel import Panel
from rich.prompt import Prompt, Confirm, IntPrompt
from rich.table import Table
from rich.theme import Theme
from rich.text import Text

try:
    import readchar
except ImportError:
    # Fallback if dependencies aren't installed yet during build/dev
    readchar = None

# ─── Theme ──────────────────────────────────────────────────────────────────────
XORFICE_THEME = Theme({
    "info":    "cyan",
    "success": "bold green",
    "warning": "bold yellow",
    "error":   "bold red",
    "accent":  "bold magenta",
    "dim":     "dim white",
    "banner":  "bold bright_cyan",
    "selected": "bold bright_red",
    "unselected": "dim white",
})

console = Console(theme=XORFICE_THEME)

# ─── Constants ──────────────────────────────────────────────────────────────────
BANNER = r"""[bold bright_cyan]
 ██╗  ██╗ ██████╗ ██████╗ ███████╗██╗ ██████╗███████╗
 ╚██╗██╔╝██╔═══██╗██╔══██╗██╔════╝██║██╔════╝██╔════╝
  ╚███╔╝ ██║   ██║██████╔╝█████╗  ██║██║     █████╗  
  ██╔██╗ ██║   ██║██╔══██╗██╔══╝  ██║██║     ██╔══╝  
 ██╔╝ ██╗╚██████╔╝██║  ██║██║     ██║╚██████╗███████╗
 ╚═╝  ╚═╝ ╚═════╝ ╚═╝  ╚═╝╚═╝     ╚═╝ ╚═════╝╚══════╝[/]
[dim]         SOTA Multimodal Inference Engine[/]
"""

VERSION = "0.1.31"
DEFAULT_HF_REPO = "Backup-bdg/Xoron-Dev-MultiMoe"
DEFAULT_HOST = "0.0.0.0"
DEFAULT_PORT = 8000

# ─── TUI Menu Logic ─────────────────────────────────────────────────────────────

class TUISelector:
    """Dynamic TUI selector with arrow-key support."""
    def __init__(self, title, options):
        self.title = title
        self.options = options # List of (label, description)
        self.index = 0
        self.running = True

    def render(self) -> Panel:
        """Render the current state of the menu."""
        content = Text()
        for i, (label, desc) in enumerate(self.options):
            is_selected = (i == self.index)
            # "Bubble" style highlighting as requested
            marker = " ● " if is_selected else " ○ "
            style = "selected" if is_selected else "unselected"
            
            line = Text()
            line.append(marker, style=style)
            line.append(f"{label:<20}", style=style)
            if desc:
                line.append(f" — {desc}", style="dim")
            
            content.append(line)
            if i < len(self.options) - 1:
                content.append("\n")

        return Panel(
            content,
            title=f"[banner]{self.title}[/]",
            border_style="bright_cyan",
            padding=(1, 3),
        )

    def run(self):
        """Run the interaction loop."""
        if readchar is None:
            console.print("[warning]readchar not installed. Falling back to numeric input.[/]")
            return self._fallback_input()

        with Live(self.render(), refresh_per_second=10, transient=True) as live:
            while self.running:
                key = readchar.readkey()
                if key == readchar.key.UP:
                    self.index = (self.index - 1) % len(self.options)
                elif key == readchar.key.DOWN:
                    self.index = (self.index + 1) % len(self.options)
                elif key == readchar.key.ENTER:
                    self.running = False
                elif key == readchar.key.CTRL_C:
                    sys.exit(0)
                
                live.update(self.render())
        
        return self.index

    def _fallback_input(self):
        """Standard input fallback."""
        console.print(self.render())
        choices = [str(i) for i in range(len(self.options))]
        val = Prompt.ask("[bold cyan]Select an option[/]", choices=choices, default="0")
        return int(val)


# ─── UI Helpers ─────────────────────────────────────────────────────────────────

def _build_info_table(host: str, port: int, model_path: str) -> Table:
    """Build a rich table summarising the server configuration."""
    table = Table(
        show_header=False,
        border_style="bright_cyan",
        padding=(0, 2),
        title="[accent]Server Configuration[/]",
    )
    table.add_column("Key", style="dim", width=20)
    table.add_column("Value", style="bold white")
    table.add_row("🌐  Endpoint",  f"http://{host}:{port}")
    table.add_row("🧠  Model",     model_path)
    table.add_row("📡  API Docs",  f"http://{host}:{port}/docs")
    table.add_row("❤️   Health",    f"http://{host}:{port}/health")
    table.add_row("🔖  Version",   VERSION)
    return table


def _build_capabilities_panel() -> Panel:
    """Build a panel showing API capabilities."""
    cap = (
        "[bold]API Endpoints[/]\n"
        "  [accent]POST[/]  /v1/chat/completions   — Text, streaming, multimodal\n"
        "  [accent]POST[/]  /v1/audio/voice-clone   — Zero-shot voice cloning\n"
        "  [accent]GET [/]  /v1/audio/voices        — List available voices\n"
        "  [accent]GET [/]  /v1/models              — List loaded models\n"
        "\n[bold]Multimodal Modes[/]\n"
        "  [dim]•[/] Text-to-Text   [dim]•[/] Speech-to-Speech   [dim]•[/] Text-to-Speech\n"
        "  [dim]•[/] Image Gen/Edit [dim]•[/] Video Gen/Edit     [dim]•[/] Agentic (tool use)\n"
        "\n[bold]Optimisations[/]\n"
        "  [dim]•[/] CUDA Graphs    [dim]•[/] Paged KV Cache     [dim]•[/] MoE Expert LRU\n"
        "  [dim]•[/] LoRA Adapters  [dim]•[/] Triton Kernels     [dim]•[/] On-the-fly Personalization"
    )
    return Panel(
        cap,
        title="[banner]Capabilities[/]",
        border_style="bright_cyan",
        padding=(1, 3),
    )


# ─── Server launcher ────────────────────────────────────────────────────────────

def _launch_server(model_path: str, host: str, port: int, run_test: bool = False, use_ssl: bool = False, domain: str = None) -> None:
    """Set env vars and launch the FastAPI server."""
    from .api import start_server  # deferred to keep CLI fast

    os.environ["XORON_MODEL_PATH"] = model_path

    console.print()
    console.rule("[accent]Inference Engine Status[/]", style="bright_cyan")
    console.print(_build_info_table(host, port, model_path))
    console.print()
    
    ssl_keyfile = None
    ssl_certfile = None
    
    if use_ssl:
        console.print("[info]🔒  Initializing Secure Context (HTTPS)...[/]")
        try:
            from .ssl_generator import generate_ssl_certs
            key_path, cert_path = generate_ssl_certs(domain=domain)
            if key_path and cert_path:
                ssl_keyfile = key_path
                ssl_certfile = cert_path
                console.print(f"[success]✔  SSL Active: {cert_path}[/]")
            else:
                console.print("[warning]⚠️  SSL Generation Failed. Falling back to HTTP.[/]")
        except Exception as e:
            console.print(f"[warning]⚠️  SSL Setup Error: {e}[/]")
            
    if run_test:
        t = threading.Thread(target=_run_auto_test, args=(host, port, bool(ssl_keyfile)), daemon=True)
        t.start()

    console.print("[info]🚀  Starting inference server…[/]\n")
    start_server(host=host, port=port, ssl_keyfile=ssl_keyfile, ssl_certfile=ssl_certfile)


def _wait_for_health(host: str, port: int, use_ssl: bool = False, timeout: int = 240) -> bool:
    """Block until /health returns 200."""
    protocol = "https" if use_ssl else "http"
    url = f"{protocol}://{host}:{port}/health"
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        try:
            req = urllib.request.Request(url)
            with urllib.request.urlopen(req, timeout=5) as resp:
                if resp.getcode() == 200:
                    return True
        except (urllib.error.URLError, OSError):
            pass
        time.sleep(2)
    return False


def _run_auto_test(host: str, port: int, use_ssl: bool = False) -> None:
    """Background: wait for server then fire a test query."""
    console.print("\n[dim]⏳  Waiting for server readiness…[/]")
    if not _wait_for_health(host, port, use_ssl):
        console.print("[error]✗  Server did not become healthy in time.[/]")
        return
    console.print()
    console.rule("[success]Auto-Test[/]", style="green")
    payload = json.dumps({
        "model": "Xoron-Dev",
        "messages": [{"role": "user", "content": "how are you?"}],
        "max_tokens": 100,
        "stream": True,
    })
    protocol = "https" if use_ssl else "http"
    subprocess.run([
        "curl", "-N", "-s", "-k", "-X", "POST",
        f"{protocol}://{host}:{port}/v1/chat/completions",
        "-H", "Content-Type: application/json",
        "-H", "x-user-id: test-user-cli",
        "-d", payload,
    ])
    console.print()
    console.rule("[success]✔  Auto-test complete — server is live[/]", style="green")
    console.print()


# ─── Interactive Menu ────────────────────────────────────────────────────────────

CONFIG_PATH = os.path.expanduser("~/.xorfice_config.json")

class ServerConfig:
    """Mutable config state for the interactive session."""
    def __init__(self):
        self.host: str = DEFAULT_HOST
        self.port: int = DEFAULT_PORT
        self.auto_test: bool = False
        self.use_ssl: bool = False
        self.domain: str = None
        self.load()

    def load(self):
        if os.path.exists(CONFIG_PATH):
            try:
                with open(CONFIG_PATH, "r") as f:
                    data = json.load(f)
                    self.host = data.get("host", self.host)
                    self.port = data.get("port", self.port)
                    self.auto_test = data.get("auto_test", self.auto_test)
                    self.use_ssl = data.get("use_ssl", self.use_ssl)
                    self.domain = data.get("domain", self.domain)
            except Exception:
                pass

    def save(self):
        try:
            with open(CONFIG_PATH, "w") as f:
                json.dump({
                    "host": self.host,
                    "port": self.port,
                    "auto_test": self.auto_test,
                    "use_ssl": self.use_ssl,
                    "domain": self.domain,
                }, f, indent=4)
        except Exception:
            pass


def _interactive_menu() -> None:
    """Main interactive TUI loop."""
    cfg = ServerConfig()

    options = [
        ("🚀  Quick Start", f"Use default model ({DEFAULT_HF_REPO})"),
        ("📦  Custom HF Model", "Download a custom model from HuggingFace"),
        ("📂  Local Weights", "Use weights already on disk"),
        ("⚙️   Server Settings", "Configure host, port, and auto-test"),
        ("ℹ️   About", "Show capabilities and version info"),
        ("退出  Exit", "Close the application"),
    ]

    while True:
        console.print()
        selector = TUISelector("Main Menu", options)
        choice_idx = selector.run()

        if choice_idx == 5: # Exit
            console.print("[dim]Goodbye! 👋[/]")
            break

        elif choice_idx == 0:
            # Quick Start
            console.print()
            load_now = Confirm.ask(
                "  [bold]Download and load the model now?[/]\n  (Select [bold red]n[/] to skip for now)",
                default=True,
            )
            os.environ["XORON_LAZY_LOAD"] = "false" if load_now else "true"
            console.print(f"\n[info]Selected:[/] [bold]{DEFAULT_HF_REPO}[/] (Lazy Load: {not load_now})")
            console.print(f"[success]✔ Loaded from Quick Start[/]")
            _launch_server(DEFAULT_HF_REPO, cfg.host, cfg.port, cfg.auto_test and load_now, cfg.use_ssl, cfg.domain)
            break

        elif choice_idx == 1:
            # Custom HF path
            console.print()
            repo_id = Prompt.ask("Enter HuggingFace Repository ID", default="Backup-bdg/Xoron-Dev-MultiMoe")
            _launch_server(repo_id, cfg.host, cfg.port, cfg.auto_test, cfg.use_ssl, cfg.domain)
            break

        elif choice_idx == 2:
            # Local weights
            console.print()
            path = Prompt.ask("Enter path to local model weights directory")
            resolved = os.path.expanduser(path)
            if not os.path.exists(resolved):
                console.print(f"[error]✗  Path not found:[/] {resolved}")
            else:
                _launch_server(resolved, cfg.host, cfg.port, cfg.auto_test, cfg.use_ssl, cfg.domain)
                break

        elif choice_idx == 3:
            # Server settings
            _settings_menu(cfg)

        elif choice_idx == 4:
            # About
            console.print()
            console.print(_build_capabilities_panel())


def _settings_menu(cfg: ServerConfig) -> None:
    """Sub-menu for server configuration."""
    console.print()
    console.rule("[accent]Server Settings[/]", style="magenta")

    cfg.host = Prompt.ask(
        "  [bold]Host[/]",
        default=cfg.host,
    )
    cfg.port = IntPrompt.ask(
        "  [bold]Port[/]",
        default=cfg.port,
    )
    cfg.auto_test = Confirm.ask(
        "  [bold]Run auto-test on startup?[/]",
        default=cfg.auto_test,
    )
    cfg.use_ssl = Confirm.ask(
        "  [bold]Generate Auto-SSL (Required for WebUI Microphones)?[/]",
        default=cfg.use_ssl,
    )
    if cfg.use_ssl:
        cfg.domain = Prompt.ask(
            "  [bold]Certbot Domain (Leave blank for Self-Signed OpenSSL)[/]",
            default=cfg.domain or "",
        )
        if not cfg.domain.strip():
            cfg.domain = None

    cfg.save()

    console.print()
    console.print("[success]✔  Settings updated and saved to ~/.xorfice_config.json[/]")
    console.print(f"  [dim]Host:[/] {cfg.host}  •  [dim]Port:[/] {cfg.port}  •  [dim]Auto-test:[/] {cfg.auto_test}  •  [dim]SSL:[/] {cfg.use_ssl}")


# ─── Entry point ─────────────────────────────────────────────────────────────────

def main() -> None:
    parser = argparse.ArgumentParser(
        prog="xorfice",
        description="xorfice — SOTA Multimodal Inference Engine",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=(
            "Examples:\n"
            "  xorfice                              # Interactive Menu (Arrow Keys)\n"
            "  xorfice --hf Backup-bdg/Xoron-Dev-MultiMoe\n"
            "  xorfice --weights ./my-weights\n"
        ),
    )
    parser.add_argument("--host",    type=str, default=DEFAULT_HOST,  help=f"Bind address  (default: {DEFAULT_HOST})")
    parser.add_argument("--port",    type=int, default=DEFAULT_PORT,  help=f"Bind port     (default: {DEFAULT_PORT})")
    parser.add_argument("--weights", type=str, default=None,          help="Path to local model weights")
    parser.add_argument("--hf",      type=str, default=None,          help="HuggingFace repo ID")
    parser.add_argument("--test",    action="store_true",             help="Auto-run a test query after startup")
    parser.add_argument("--ssl",     action="store_true",             help="Auto-generate SSL Certificates for HTTPS mode")
    parser.add_argument("--domain",  type=str, default=None,          help="Optional domain to generate Certbot certificates instead of local Self-Signed")
    parser.add_argument("--version", action="version",                version=f"xorfice {VERSION}")

    args = parser.parse_args()

    # ── Show banner always ──
    console.print(BANNER)

    # ── Direct launch if model specified via flags ──
    if args.hf:
        _launch_server(args.hf, args.host, args.port, args.test, args.ssl, args.domain)
    elif args.weights:
        resolved = os.path.expanduser(args.weights)
        if not os.path.exists(resolved):
            console.print(f"[error]✗  Local path not found:[/] {resolved}")
            sys.exit(1)
        _launch_server(resolved, args.host, args.port, args.test, args.ssl, args.domain)
    else:
        # ── Interactive TUI Menu ──
        _interactive_menu()


if __name__ == "__main__":
    main()
