#!/usr/bin/env python3
"""
ICC AI/HPC Stack — GUI Installer
Copyright (c) 2026 International Computer Concepts (ICC)
https://icc-usa.com

A graphical installer for the ICC AI/HPC Stack.
Supports both NVIDIA (CUDA) and AMD (ROCm) GPU platforms.

Requirements: Python 3.10+ with tkinter (pre-installed on Ubuntu 24.04)

Usage:
    sudo python3 icc-ai-stack-gui.py
    # or
    chmod +x icc-ai-stack-gui.py && sudo ./icc-ai-stack-gui.py
"""

import tkinter as tk
from tkinter import ttk, scrolledtext, messagebox
import subprocess
import threading
import os
import sys
import shutil
from pathlib import Path

# ─────────────────────────────────────────────────────────────────────────────
# Constants
# ─────────────────────────────────────────────────────────────────────────────
APP_VERSION = "2.0.0"
APP_TITLE = "ICC AI/HPC Stack Installer"

# ICC brand colors
ICC_BG = "#0a0a0a"
ICC_BG2 = "#111111"
ICC_BG3 = "#161616"
ICC_CARD = "#141414"
ICC_BORDER = "#2a2a2a"
ICC_TEXT = "#ffffff"
ICC_TEXT2 = "#a0a0a0"
ICC_TEXT3 = "#666666"
ICC_GREEN = "#76b900"
ICC_GREEN_HOVER = "#8ad400"
ICC_CYAN = "#22d3ee"
ICC_RED = "#ef4444"
ICC_AMD_RED = "#ed1c24"
ICC_ORANGE = "#f97316"

NVIDIA_VERSIONS = {
    "cuda": ["12.8", "12.6", "12.4"],
    "driver": ["570", "560", "550", "535"],
}

AMD_VERSIONS = {
    "rocm": ["6.3", "6.2", "6.1"],
}


class ICCInstallerApp:
    def __init__(self, root):
        self.root = root
        self.root.title(APP_TITLE)
        self.root.geometry("960x780")
        self.root.minsize(800, 700)
        self.root.configure(bg=ICC_BG)

        # State
        self.gpu_vendor = tk.StringVar(value="nvidia")
        self.cuda_version = tk.StringVar(value="12.8")
        self.driver_version = tk.StringVar(value="570")
        self.rocm_version = tk.StringVar(value="6.3")
        self.install_frameworks = tk.BooleanVar(value=True)
        self.install_jupyter = tk.BooleanVar(value=True)
        self.with_docker = tk.BooleanVar(value=False)
        self.with_profiling = tk.BooleanVar(value=False)
        self.with_inference = tk.BooleanVar(value=False)
        self.with_cluster = tk.BooleanVar(value=False)
        self.dry_run = tk.BooleanVar(value=False)
        self.is_installing = False
        self.process = None

        self._setup_styles()
        self._build_ui()
        self._detect_gpu()

    # ─── Styles ───────────────────────────────────────────────────────────
    def _setup_styles(self):
        self.style = ttk.Style()
        self.style.theme_use("clam")

        self.style.configure(".", background=ICC_BG, foreground=ICC_TEXT,
                             fieldbackground=ICC_BG2, borderwidth=0)
        self.style.configure("TFrame", background=ICC_BG)
        self.style.configure("Card.TFrame", background=ICC_CARD)
        self.style.configure("TLabel", background=ICC_BG, foreground=ICC_TEXT,
                             font=("Inter", 11))
        self.style.configure("Header.TLabel", font=("Inter", 24, "bold"),
                             foreground=ICC_TEXT)
        self.style.configure("Sub.TLabel", font=("Inter", 12),
                             foreground=ICC_TEXT2)
        self.style.configure("Section.TLabel", font=("Inter", 13, "bold"),
                             foreground=ICC_TEXT)
        self.style.configure("Muted.TLabel", foreground=ICC_TEXT3,
                             font=("Inter", 10))
        self.style.configure("Green.TLabel", foreground=ICC_GREEN,
                             font=("Inter", 11, "bold"))
        self.style.configure("TCheckbutton", background=ICC_BG,
                             foreground=ICC_TEXT, font=("Inter", 11))
        self.style.configure("TRadiobutton", background=ICC_BG,
                             foreground=ICC_TEXT, font=("Inter", 11))
        self.style.configure("TCombobox", fieldbackground=ICC_BG2,
                             background=ICC_BG3, foreground=ICC_TEXT,
                             arrowcolor=ICC_TEXT2)
        self.style.map("TCombobox",
                       fieldbackground=[("readonly", ICC_BG2)],
                       foreground=[("readonly", ICC_TEXT)])
        self.style.configure("TNotebook", background=ICC_BG, borderwidth=0)
        self.style.configure("TNotebook.Tab", background=ICC_BG2,
                             foreground=ICC_TEXT2, padding=[16, 8],
                             font=("Inter", 11))
        self.style.map("TNotebook.Tab",
                       background=[("selected", ICC_BG)],
                       foreground=[("selected", ICC_TEXT)])

        self.style.configure("Green.Horizontal.TProgressbar",
                             troughcolor=ICC_BG2, background=ICC_GREEN,
                             darkcolor=ICC_GREEN, lightcolor=ICC_GREEN,
                             bordercolor=ICC_BORDER)

    # ─── UI Construction ──────────────────────────────────────────────────
    def _build_ui(self):
        main = ttk.Frame(self.root, padding=32)
        main.pack(fill="both", expand=True)

        header_frame = ttk.Frame(main)
        header_frame.pack(fill="x", pady=(0, 24))

        ttk.Label(header_frame, text="ICC AI/HPC Stack",
                  style="Header.TLabel").pack(side="left")
        ttk.Label(header_frame, text=f"v{APP_VERSION}",
                  style="Muted.TLabel").pack(side="left", padx=(12, 0),
                                              pady=(8, 0))

        self.gpu_info_label = ttk.Label(header_frame, text="Detecting GPU...",
                                         style="Sub.TLabel")
        self.gpu_info_label.pack(side="right")

        self.notebook = ttk.Notebook(main)
        self.notebook.pack(fill="both", expand=True, pady=(0, 16))

        self._build_config_tab()
        self._build_log_tab()
        self._build_bottom_bar(main)

    def _build_config_tab(self):
        config_frame = ttk.Frame(self.notebook, padding=16)
        self.notebook.add(config_frame, text="  Configuration  ")

        canvas = tk.Canvas(config_frame, bg=ICC_BG, highlightthickness=0)
        scrollbar = ttk.Scrollbar(config_frame, orient="vertical",
                                   command=canvas.yview)
        scroll_frame = ttk.Frame(canvas)

        scroll_frame.bind("<Configure>",
                          lambda e: canvas.configure(
                              scrollregion=canvas.bbox("all")))

        canvas.create_window((0, 0), window=scroll_frame, anchor="nw")
        canvas.configure(yscrollcommand=scrollbar.set)

        canvas.pack(side="left", fill="both", expand=True)
        scrollbar.pack(side="right", fill="y")

        def _on_mousewheel(event):
            canvas.yview_scroll(int(-1 * (event.delta / 120)), "units")
        canvas.bind_all("<MouseWheel>", _on_mousewheel)
        canvas.bind_all("<Button-4>", lambda e: canvas.yview_scroll(-1, "units"))
        canvas.bind_all("<Button-5>", lambda e: canvas.yview_scroll(1, "units"))

        # GPU Vendor
        self._section_header(scroll_frame, "GPU Platform")
        vendor_frame = ttk.Frame(scroll_frame)
        vendor_frame.pack(fill="x", padx=8, pady=(0, 16))
        ttk.Radiobutton(vendor_frame, text="NVIDIA (CUDA)",
                         variable=self.gpu_vendor, value="nvidia",
                         command=self._on_vendor_change).pack(side="left",
                                                                padx=(0, 24))
        ttk.Radiobutton(vendor_frame, text="AMD (ROCm)",
                         variable=self.gpu_vendor, value="amd",
                         command=self._on_vendor_change).pack(side="left")

        # Versions
        self._section_header(scroll_frame, "Versions")
        ver_frame = ttk.Frame(scroll_frame)
        ver_frame.pack(fill="x", padx=8, pady=(0, 16))

        self.nvidia_ver_frame = ttk.Frame(ver_frame)
        self.nvidia_ver_frame.pack(fill="x")
        row1 = ttk.Frame(self.nvidia_ver_frame)
        row1.pack(fill="x", pady=4)
        ttk.Label(row1, text="CUDA Version:", width=16).pack(side="left")
        ttk.Combobox(row1, textvariable=self.cuda_version,
                     values=NVIDIA_VERSIONS["cuda"],
                     state="readonly", width=12).pack(side="left", padx=(0, 24))
        ttk.Label(row1, text="Driver Branch:", width=14).pack(side="left")
        ttk.Combobox(row1, textvariable=self.driver_version,
                     values=NVIDIA_VERSIONS["driver"],
                     state="readonly", width=12).pack(side="left")

        self.amd_ver_frame = ttk.Frame(ver_frame)
        amd_row = ttk.Frame(self.amd_ver_frame)
        amd_row.pack(fill="x", pady=4)
        ttk.Label(amd_row, text="ROCm Version:", width=16).pack(side="left")
        ttk.Combobox(amd_row, textvariable=self.rocm_version,
                     values=AMD_VERSIONS["rocm"],
                     state="readonly", width=12).pack(side="left", padx=(0, 24))
        ttk.Label(amd_row, text="Target: MI300 (gfx942)",
                  style="Muted.TLabel").pack(side="left")

        # Core Components
        self._section_header(scroll_frame, "Core Components")
        core_frame = ttk.Frame(scroll_frame)
        core_frame.pack(fill="x", padx=8, pady=(0, 16))
        ttk.Checkbutton(core_frame, text="Deep Learning Frameworks (PyTorch, TensorFlow, JAX)",
                         variable=self.install_frameworks).pack(anchor="w", pady=2)
        ttk.Checkbutton(core_frame, text="JupyterLab",
                         variable=self.install_jupyter).pack(anchor="w", pady=2)

        # Optional Modules
        self._section_header(scroll_frame, "Optional Modules")
        mod_frame = ttk.Frame(scroll_frame)
        mod_frame.pack(fill="x", padx=8, pady=(0, 16))
        modules = [
            (self.with_docker, "Docker + GPU Container Runtime",
             "Docker CE with GPU passthrough (NVIDIA CTK or ROCm device mount)"),
            (self.with_profiling, "GPU Profiling & Debugging",
             "NVIDIA: DCGM + Nsight  |  AMD: rocprof + omniperf + omnitrace"),
            (self.with_inference, "Inference Stack",
             "vLLM, Triton client, llama.cpp (CUDA or HIP build)"),
            (self.with_cluster, "Cluster & Multi-Node",
             "OpenMPI, NCCL/RCCL tests, bandwidth test, pssh"),
        ]
        for var, label, desc in modules:
            f = ttk.Frame(mod_frame)
            f.pack(fill="x", pady=4)
            ttk.Checkbutton(f, text=label, variable=var).pack(anchor="w")
            ttk.Label(f, text=desc, style="Muted.TLabel").pack(anchor="w",
                                                                 padx=(28, 0))

        # Options
        self._section_header(scroll_frame, "Options")
        opt_frame = ttk.Frame(scroll_frame)
        opt_frame.pack(fill="x", padx=8, pady=(0, 16))
        ttk.Checkbutton(opt_frame, text="Dry Run (preview only — no changes)",
                         variable=self.dry_run).pack(anchor="w", pady=2)

        # Command Preview
        self._section_header(scroll_frame, "Command Preview")
        self.cmd_preview = tk.Text(scroll_frame, height=3, bg=ICC_BG2,
                                    fg=ICC_CYAN, font=("JetBrains Mono", 11),
                                    relief="flat", padx=16, pady=12,
                                    wrap="word", state="disabled",
                                    highlightthickness=1,
                                    highlightcolor=ICC_BORDER,
                                    highlightbackground=ICC_BORDER)
        self.cmd_preview.pack(fill="x", padx=8, pady=(0, 8))

        for var in [self.gpu_vendor, self.cuda_version, self.driver_version,
                    self.rocm_version, self.install_frameworks,
                    self.install_jupyter, self.with_docker,
                    self.with_profiling, self.with_inference,
                    self.with_cluster, self.dry_run]:
            var.trace_add("write", lambda *_: self._update_preview())

        self._update_preview()

    def _build_log_tab(self):
        log_frame = ttk.Frame(self.notebook, padding=16)
        self.notebook.add(log_frame, text="  Installation Log  ")

        self.log_text = scrolledtext.ScrolledText(
            log_frame, bg=ICC_BG2, fg=ICC_TEXT2,
            font=("JetBrains Mono", 10), relief="flat",
            padx=16, pady=12, state="disabled",
            insertbackground=ICC_TEXT,
            highlightthickness=1,
            highlightcolor=ICC_BORDER,
            highlightbackground=ICC_BORDER,
        )
        self.log_text.pack(fill="both", expand=True)

        self.log_text.tag_configure("phase", foreground=ICC_GREEN)
        self.log_text.tag_configure("warn", foreground=ICC_ORANGE)
        self.log_text.tag_configure("error", foreground=ICC_RED)
        self.log_text.tag_configure("success", foreground="#28c840")
        self.log_text.tag_configure("info", foreground=ICC_CYAN)

    def _build_bottom_bar(self, parent):
        bar = ttk.Frame(parent)
        bar.pack(fill="x", pady=(0, 0))

        self.progress = ttk.Progressbar(bar, mode="indeterminate",
                                         style="Green.Horizontal.TProgressbar")
        self.progress.pack(fill="x", pady=(0, 12))

        btn_frame = ttk.Frame(bar)
        btn_frame.pack(fill="x")

        self.status_label = ttk.Label(btn_frame, text="Ready",
                                       style="Muted.TLabel")
        self.status_label.pack(side="left")

        self.install_btn = tk.Button(
            btn_frame, text="Install", font=("Inter", 13, "bold"),
            bg=ICC_GREEN, fg=ICC_BG, activebackground=ICC_GREEN_HOVER,
            activeforeground=ICC_BG, relief="flat", padx=32, pady=10,
            cursor="hand2", command=self._on_install,
        )
        self.install_btn.pack(side="right")

        self.cancel_btn = tk.Button(
            btn_frame, text="Cancel", font=("Inter", 11),
            bg=ICC_BG3, fg=ICC_TEXT2, activebackground=ICC_BORDER,
            activeforeground=ICC_TEXT, relief="flat", padx=20, pady=10,
            cursor="hand2", command=self._on_cancel, state="disabled",
        )
        self.cancel_btn.pack(side="right", padx=(0, 8))

    # ─── Helpers ──────────────────────────────────────────────────────────
    def _section_header(self, parent, text):
        f = ttk.Frame(parent)
        f.pack(fill="x", pady=(16, 8))
        ttk.Label(f, text=text, style="Section.TLabel").pack(anchor="w")
        ttk.Separator(f, orient="horizontal").pack(fill="x", pady=(6, 0))

    def _on_vendor_change(self):
        if self.gpu_vendor.get() == "nvidia":
            self.amd_ver_frame.pack_forget()
            self.nvidia_ver_frame.pack(fill="x")
        else:
            self.nvidia_ver_frame.pack_forget()
            self.amd_ver_frame.pack(fill="x")
        self._update_preview()

    def _detect_gpu(self):
        def _detect():
            try:
                result = subprocess.run(
                    ["lspci"], capture_output=True, text=True, timeout=5
                )
                output = result.stdout.lower()
                if "nvidia" in output:
                    gpu_lines = [l for l in result.stdout.splitlines()
                                 if "nvidia" in l.lower()]
                    name = gpu_lines[0].split(":")[-1].strip() if gpu_lines else "NVIDIA GPU"
                    self.root.after(0, lambda: self.gpu_info_label.configure(
                        text=f"Detected: {name[:50]}"))
                    self.root.after(0, lambda: self.gpu_vendor.set("nvidia"))
                elif "amd" in output and ("instinct" in output or "mi3" in output):
                    gpu_lines = [l for l in result.stdout.splitlines()
                                 if "amd" in l.lower() and ("instinct" in l.lower() or "mi3" in l.lower())]
                    name = gpu_lines[0].split(":")[-1].strip() if gpu_lines else "AMD Instinct"
                    self.root.after(0, lambda: self.gpu_info_label.configure(
                        text=f"Detected: {name[:50]}"))
                    self.root.after(0, lambda: self.gpu_vendor.set("amd"))
                    self.root.after(0, self._on_vendor_change)
                else:
                    self.root.after(0, lambda: self.gpu_info_label.configure(
                        text="No supported GPU detected"))
            except Exception:
                self.root.after(0, lambda: self.gpu_info_label.configure(
                    text="GPU detection unavailable"))

        threading.Thread(target=_detect, daemon=True).start()

    def _build_command(self):
        vendor = self.gpu_vendor.get()
        if vendor == "nvidia":
            script = "./install-icc-ai-stack.sh"
            cmd = [script]
            cmd += ["--cuda-version", self.cuda_version.get()]
            cmd += ["--driver-version", self.driver_version.get()]
        else:
            script = "./install-icc-ai-stack-amd.sh"
            cmd = [script]
            cmd += ["--rocm-version", self.rocm_version.get()]

        if not self.install_frameworks.get():
            cmd.append("--no-frameworks")
        if not self.install_jupyter.get():
            cmd.append("--no-jupyter")
        if self.with_docker.get():
            cmd.append("--with-docker")
        if self.with_profiling.get():
            cmd.append("--with-profiling")
        if self.with_inference.get():
            cmd.append("--with-inference")
        if self.with_cluster.get():
            cmd.append("--with-cluster")
        if self.dry_run.get():
            cmd.append("--dry-run")

        cmd.append("--yes")
        return cmd

    def _update_preview(self, *_):
        cmd = self._build_command()
        self.cmd_preview.configure(state="normal")
        self.cmd_preview.delete("1.0", "end")
        self.cmd_preview.insert("1.0", "sudo " + " ".join(cmd))
        self.cmd_preview.configure(state="disabled")

    def _log_append(self, text, tag=None):
        self.log_text.configure(state="normal")
        if tag:
            self.log_text.insert("end", text + "\n", tag)
        else:
            self.log_text.insert("end", text + "\n")
        self.log_text.see("end")
        self.log_text.configure(state="disabled")

    # ─── Installation ───────────────────────────────────────────────────
    def _on_install(self):
        if self.is_installing:
            return

        if os.geteuid() != 0 and not self.dry_run.get():
            messagebox.showwarning(
                "Root Required",
                "This installer must be run as root.\n\n"
                "Run: sudo python3 icc-ai-stack-gui.py\n\n"
                "Or enable 'Dry Run' to preview without changes."
            )
            return

        cmd = self._build_command()
        script = cmd[0]

        search_paths = [
            Path(script),
            Path(__file__).parent / script.lstrip("./"),
            Path("/usr/local/bin") / script.lstrip("./"),
            Path.home() / script.lstrip("./"),
        ]

        found_script = None
        for p in search_paths:
            if p.exists():
                found_script = str(p)
                break

        if not found_script:
            messagebox.showerror(
                "Script Not Found",
                f"Cannot find {script}\n\n"
                f"Place the script in the same directory as this GUI,\n"
                f"or in /usr/local/bin/."
            )
            return

        cmd[0] = found_script

        self.is_installing = True
        self.install_btn.configure(state="disabled", bg=ICC_TEXT3)
        self.cancel_btn.configure(state="normal")
        self.progress.start(10)
        self.status_label.configure(text="Installing...")
        self.notebook.select(1)

        self._log_append(f"{'─' * 60}", "info")
        self._log_append(f"ICC AI/HPC Stack Installer — Starting", "info")
        self._log_append(f"Command: sudo {' '.join(cmd)}", "info")
        self._log_append(f"{'─' * 60}", "info")
        self._log_append("")

        threading.Thread(target=self._run_install, args=(cmd,),
                         daemon=True).start()

    def _run_install(self, cmd):
        try:
            self.process = subprocess.Popen(
                cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                text=True, bufsize=1, universal_newlines=True,
            )

            for line in self.process.stdout:
                line = line.rstrip()
                if not line:
                    continue

                tag = None
                if "[ICC]" in line:
                    tag = "phase"
                elif "[WARN]" in line:
                    tag = "warn"
                elif "[ERROR]" in line:
                    tag = "error"
                elif "\u2713" in line:
                    tag = "success"
                elif "[INFO]" in line:
                    tag = "info"

                if "Phase " in line and "/" in line:
                    self.root.after(0, lambda l=line: self.status_label.configure(
                        text=l.split("]", 1)[-1].strip()[:80]))

                self.root.after(0, lambda l=line, t=tag: self._log_append(l, t))

            self.process.wait()
            rc = self.process.returncode
            self.root.after(0, lambda: self._install_complete(rc))

        except Exception as e:
            self.root.after(0, lambda: self._log_append(
                f"Error: {str(e)}", "error"))
            self.root.after(0, lambda: self._install_complete(1))

    def _install_complete(self, return_code):
        self.is_installing = False
        self.process = None
        self.progress.stop()
        self.cancel_btn.configure(state="disabled")
        self.install_btn.configure(state="normal", bg=ICC_GREEN)

        if return_code == 0:
            self.status_label.configure(text="Installation complete!")
            self._log_append("")
            self._log_append("\u2500" * 60, "success")
            self._log_append("Installation completed successfully!", "success")
            self._log_append("A reboot is recommended.", "warn")
            self._log_append("\u2500" * 60, "success")
            messagebox.showinfo("Complete",
                                "ICC AI/HPC Stack installed successfully!\n\n"
                                "A reboot is recommended.")
        else:
            self.status_label.configure(text=f"Installation failed (exit {return_code})")
            self._log_append("")
            self._log_append(f"Installation failed with exit code {return_code}", "error")
            messagebox.showerror("Failed",
                                 f"Installation failed (exit code {return_code}).\n\n"
                                 "Check the log tab for details.")

    def _on_cancel(self):
        if self.process and self.is_installing:
            if messagebox.askyesno("Cancel",
                                    "Cancel the installation?\n\n"
                                    "The system may be in a partial state."):
                try:
                    self.process.terminate()
                except Exception:
                    pass
                self._log_append("Installation cancelled by user.", "warn")
                self._install_complete(1)


# ─────────────────────────────────────────────────────────────────────────────
# Entry point
# ─────────────────────────────────────────────────────────────────────────────
def main():
    root = tk.Tk()
    try:
        root.iconname("ICC AI Stack")
    except Exception:
        pass
    app = ICCInstallerApp(root)
    root.mainloop()


if __name__ == "__main__":
    main()