#!/usr/bin/env python3
"""
VelionSoft Ubuntu AI Server Installer.

Installs Docker, Ollama, Open WebUI and optionally pulls selected Ollama models.
Target OS: Ubuntu/Debian servers with systemd and apt.
"""

from __future__ import annotations

import os
import re
import shutil
import socket
import subprocess
import sys
import time
from pathlib import Path
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen


OLLAMA_PORT = 11434
WEBUI_PORT = 3000
OPEN_WEBUI_CONTAINER = "open-webui"
OPEN_WEBUI_IMAGE = "ghcr.io/open-webui/open-webui:main"
OLLAMA_LIBRARY_URL = "https://ollama.com/library"

RECOMMENDED_MODELS = [
    "qwen2.5-coder:7b",
    "qwen2.5-coder:14b",
    "qwen2.5-coder:32b",
    "deepseek-coder-v2:16b",
    "codestral:22b",
    "llama3.1:8b",
    "mistral:7b",
    "gemma3:12b",
]


def run(command: list[str], check: bool = True) -> int:
    print("\n>>> " + " ".join(command))
    result = subprocess.run(command, text=True, check=False)
    if check and result.returncode != 0:
        print("\nERROR: Command failed:")
        print(" ".join(command))
        sys.exit(result.returncode)
    return result.returncode


def run_shell(command: str, check: bool = True) -> int:
    print(f"\n>>> {command}")
    result = subprocess.run(["bash", "-lc", command], text=True, check=False)
    if check and result.returncode != 0:
        print("\nERROR: Command failed:")
        print(command)
        sys.exit(result.returncode)
    return result.returncode


def command_exists(command: str) -> bool:
    return shutil.which(command) is not None


def require_root() -> None:
    if os.geteuid() != 0:
        print("ERROR: run this script as root.")
        print("Example:")
        print("sudo python3 install_ai_server.py")
        sys.exit(1)


def get_public_ip() -> str:
    services = [
        "https://api.ipify.org",
        "https://ifconfig.me",
        "https://icanhazip.com",
    ]

    for service in services:
        try:
            result = subprocess.check_output(
                ["curl", "-4", "-s", "--max-time", "5", service],
                text=True,
            ).strip()
            if re.match(r"^\d{1,3}(\.\d{1,3}){3}$", result):
                return result
        except Exception:
            pass

    try:
        hostname = socket.gethostname()
        ip = socket.gethostbyname(hostname)
        if ip:
            return ip
    except Exception:
        pass

    return "SERVER_IP"


def install_base_packages() -> None:
    run(["apt", "update"])
    run(["apt", "install", "-y", "curl", "ca-certificates", "gnupg", "lsb-release", "ufw", "python3"])


def install_ollama() -> None:
    if command_exists("ollama"):
        print("\n>>> Ollama is already installed.")
    else:
        run_shell("curl -fsSL https://ollama.com/install.sh | sh")

    override_dir = Path("/etc/systemd/system/ollama.service.d")
    override_dir.mkdir(parents=True, exist_ok=True)

    override_file = override_dir / "override.conf"
    override_file.write_text(
        f"""[Service]
Environment="OLLAMA_HOST=0.0.0.0:{OLLAMA_PORT}"
""",
        encoding="utf-8",
    )

    run(["systemctl", "daemon-reload"])
    run(["systemctl", "enable", "--now", "ollama"])
    run(["systemctl", "restart", "ollama"])


def install_docker() -> None:
    if command_exists("docker"):
        print("\n>>> Docker is already installed.")
    else:
        run(["apt", "install", "-y", "docker.io", "docker-compose"])

    run(["systemctl", "enable", "--now", "docker"])


def setup_firewall() -> None:
    if command_exists("ufw"):
        run(["ufw", "allow", f"{WEBUI_PORT}/tcp"], check=False)
        run(["ufw", "allow", f"{OLLAMA_PORT}/tcp"], check=False)


def run_open_webui() -> None:
    run(["docker", "rm", "-f", OPEN_WEBUI_CONTAINER], check=False)
    run(
        [
            "docker",
            "run",
            "-d",
            "-p",
            f"{WEBUI_PORT}:8080",
            "--add-host=host.docker.internal:host-gateway",
            "-e",
            f"OLLAMA_BASE_URL=http://host.docker.internal:{OLLAMA_PORT}",
            "-v",
            "open-webui:/app/backend/data",
            "--name",
            OPEN_WEBUI_CONTAINER,
            "--restart",
            "always",
            OPEN_WEBUI_IMAGE,
        ]
    )
    print("\n>>> Waiting for Open WebUI to start...")
    time.sleep(10)


def fetch_ollama_models() -> list[str]:
    print("\n>>> Fetching Ollama Library model list...")
    try:
        request = Request(OLLAMA_LIBRARY_URL, headers={"User-Agent": "Mozilla/5.0"})
        with urlopen(request, timeout=20) as response:
            html = response.read().decode("utf-8", errors="ignore")

        found = re.findall(r'href="/library/([a-zA-Z0-9._-]+)"', html)
        return sorted(set(found))
    except (HTTPError, URLError, TimeoutError) as error:
        print(f"\nWARNING: could not fetch online model list: {error}")
        return []


def filter_models(models: list[str]) -> list[str]:
    while True:
        keyword = input("\nModel filter, e.g. qwen, deepseek, llama, gemma or ENTER for full list: ").strip().lower()
        if keyword == "":
            return models
        filtered = [model for model in models if keyword in model.lower()]
        if filtered:
            return filtered
        print("No results for that filter.")


def choose_models(online_models: list[str]) -> list[str]:
    print("\nHow do you want to choose models?")
    print("1. Recommended list")
    print("2. Online Ollama Library list")
    print("3. Type model/tag manually")
    print("4. Skip model installation")

    mode = input("\nChoice [1-4]: ").strip()
    if mode == "4":
        return []

    if mode == "3":
        manual = input("\nEnter comma-separated models, e.g. qwen2.5-coder:14b,codestral:22b: ").strip()
        return [item.strip() for item in manual.split(",") if item.strip()]

    if mode == "2":
        selected_source = filter_models(online_models) if online_models else RECOMMENDED_MODELS
    else:
        selected_source = RECOMMENDED_MODELS

    print("\nAvailable models:")
    for index, model in enumerate(selected_source, start=1):
        print(f"{index}. {model}")

    print("\nExamples:")
    print("1       = install one model")
    print("1,2,4   = install multiple models")
    print("all     = install all shown models")
    print("ENTER   = skip model installation")

    choice = input("\nChoose models: ").strip().lower()
    if choice == "":
        return []
    if choice == "all":
        return selected_source

    selected = []
    for part in choice.split(","):
        part = part.strip()
        if not part.isdigit():
            print(f"Skipping unknown input: {part}")
            continue
        index = int(part)
        if 1 <= index <= len(selected_source):
            selected.append(selected_source[index - 1])
        else:
            print(f"Skipping invalid number: {index}")

    return selected


def pull_models(models: list[str]) -> None:
    if not models:
        print("\n>>> No model selected.")
        return

    print("\nSelected models:")
    for model in models:
        print(f"- {model}")

    confirm = input("\nInstall these models? [y/N]: ").strip().lower()
    if confirm != "y":
        print("Model installation skipped.")
        return

    for model in models:
        run(["ollama", "pull", model], check=False)


def verify_services() -> None:
    print("\n>>> Service check")
    run(["systemctl", "is-active", "ollama"], check=False)
    run(["docker", "ps"], check=False)
    run(["curl", "-s", f"http://127.0.0.1:{OLLAMA_PORT}/api/tags"], check=False)


def main() -> None:
    require_root()

    print("\n========================================")
    print("  VelionSoft Ubuntu AI Server Installer")
    print("========================================")

    install_base_packages()
    install_ollama()
    install_docker()
    setup_firewall()
    run_open_webui()

    online_models = fetch_ollama_models()
    selected_models = choose_models(online_models)
    pull_models(selected_models)

    verify_services()
    public_ip = get_public_ip()

    print("\n========================================")
    print("  INSTALLATION COMPLETE")
    print("========================================")
    print("Open WebUI:")
    print(f"http://{public_ip}:{WEBUI_PORT}")
    print("")
    print("Ollama API:")
    print(f"http://{public_ip}:{OLLAMA_PORT}")
    print("")
    print("IMPORTANT:")
    print("Port 11434 is the Ollama API. Do not leave it publicly exposed without protection.")
    print("For production, restrict firewall access to your IP or use an Nginx reverse proxy.")
    print("========================================\n")


if __name__ == "__main__":
    main()