#!/usr/bin/env python3
"""
Bootstrap a fresh Debian headless install into a CasaOS-based home NAS.

Target use:
  1. Install Debian minimal/headless manually onto the OS drive only.
  2. Boot Debian, log in as your admin user.
  3. Copy this script onto the machine.
  4. Run: sudo python3 bootstrap_casaos_nas.py --hostname dell-nas --admin-user <your-user>

This script intentionally DOES NOT partition, format, or mount media drives. Media drives
hold real data, so handle them separately after verifying disk IDs/UUIDs.
"""

from __future__ import annotations

import argparse
import datetime as dt
import os
import platform
import re
import shutil
import subprocess
import sys
from pathlib import Path

LOG_DIR = Path("/var/log/hephaestus-nas-bootstrap")
LOG_FILE = LOG_DIR / "bootstrap.log"
BACKUP_DIR = Path("/root/hephaestus-nas-bootstrap-backups")

APT_PACKAGES = [
    # Base admin tools
    "ca-certificates",
    "curl",
    "wget",
    "gnupg",
    "lsb-release",
    "apt-transport-https",
    "software-properties-common",
    "sudo",
    "nano",
    "vim-tiny",
    "htop",
    "tmux",
    "git",
    "jq",
    "unzip",
    "bash-completion",
    # Network and discovery
    "openssh-server",
    "avahi-daemon",
    "net-tools",
    "iproute2",
    "dnsutils",
    "nmap",
    # NAS/filesystem tools
    "samba",
    "smbclient",
    "cifs-utils",
    "rsync",
    "acl",
    "attr",
    "xfsprogs",
    "btrfs-progs",
    "exfatprogs",
    "ntfs-3g",
    # Disk health and monitoring
    "smartmontools",
    "hdparm",
    "nvme-cli",
    "lm-sensors",
    "iotop",
    # Safety/security
    "ufw",
    "fail2ban",
    "unattended-upgrades",
    "logrotate",
    # Useful web admin panel outside CasaOS
    "cockpit",
]

SMB_MANAGED_BEGIN = "# BEGIN HEPHAESTUS MANAGED NAS SHARES"
SMB_MANAGED_END = "# END HEPHAESTUS MANAGED NAS SHARES"

DEFAULT_SHARE_BLOCK = f"""
{SMB_MANAGED_BEGIN}
[media_primary]
   path = /mnt/media_primary
   browseable = yes
   read only = no
   guest ok = no
   valid users = @sambashare
   create mask = 0664
   directory mask = 0775
   force group = sambashare

[media_backup]
   path = /mnt/media_backup
   browseable = yes
   read only = no
   guest ok = no
   valid users = @sambashare
   create mask = 0664
   directory mask = 0775
   force group = sambashare
{SMB_MANAGED_END}
""".strip() + "\n"


def log(message: str) -> None:
    timestamp = dt.datetime.now(dt.timezone.utc).astimezone().isoformat(timespec="seconds")
    line = f"[{timestamp}] {message}"
    print(line)
    try:
        LOG_DIR.mkdir(parents=True, exist_ok=True)
        with LOG_FILE.open("a", encoding="utf-8") as fh:
            fh.write(line + "\n")
    except PermissionError:
        pass


def run(cmd: list[str], *, check: bool = True, env: dict[str, str] | None = None) -> subprocess.CompletedProcess[str]:
    log("RUN: " + " ".join(sh_quote(x) for x in cmd))
    merged_env = os.environ.copy()
    if env:
        merged_env.update(env)
    proc = subprocess.run(cmd, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, env=merged_env)
    if proc.stdout.strip():
        for line in proc.stdout.rstrip().splitlines():
            log("  " + line)
    if check and proc.returncode != 0:
        raise RuntimeError(f"Command failed with exit {proc.returncode}: {' '.join(cmd)}")
    return proc


def sh_quote(value: str) -> str:
    if re.fullmatch(r"[A-Za-z0-9_./:=@%+-]+", value):
        return value
    return "'" + value.replace("'", "'\\''") + "'"


def require_root() -> None:
    if os.geteuid() != 0:
        raise SystemExit("Run this as root, e.g. sudo python3 bootstrap_casaos_nas.py")


def check_debian_family() -> None:
    os_release = Path("/etc/os-release").read_text(encoding="utf-8", errors="ignore") if Path("/etc/os-release").exists() else ""
    if "debian" not in os_release.lower() and "ubuntu" not in os_release.lower():
        raise SystemExit("This script is intended for Debian-family systems only.")
    log("OS detected from /etc/os-release: " + " | ".join(os_release.splitlines()[:4]))
    log(f"Kernel/arch: {platform.release()} / {platform.machine()}")


def backup_file(path: Path) -> Path | None:
    if not path.exists():
        return None
    BACKUP_DIR.mkdir(parents=True, exist_ok=True)
    stamp = dt.datetime.now().strftime("%Y%m%d-%H%M%S")
    dest = BACKUP_DIR / f"{path.name}.{stamp}.bak"
    shutil.copy2(path, dest)
    log(f"Backed up {path} -> {dest}")
    return dest


def set_hostname(hostname: str | None) -> None:
    if not hostname:
        return
    if not re.fullmatch(r"[a-zA-Z0-9][a-zA-Z0-9-]{0,62}", hostname):
        raise SystemExit("Invalid hostname. Use letters, numbers, hyphens; no spaces.")
    current = platform.node()
    if current == hostname:
        log(f"Hostname already set to {hostname}")
        return
    run(["hostnamectl", "set-hostname", hostname])
    log(f"Hostname changed from {current!r} to {hostname!r}; reboot recommended later.")


def apt_install() -> None:
    env = {"DEBIAN_FRONTEND": "noninteractive"}
    run(["apt-get", "update"], env=env)
    run(["apt-get", "upgrade", "-y"], env=env)
    run(["apt-get", "install", "-y", *APT_PACKAGES], env=env)


def enable_services() -> None:
    services = [
        "ssh",
        "smbd",
        "nmbd",
        "avahi-daemon",
        "smartmontools",
        "fail2ban",
        "cockpit.socket",
    ]
    for svc in services:
        run(["systemctl", "enable", "--now", svc], check=False)
    run(["systemctl", "status", "ssh", "--no-pager"], check=False)
    run(["systemctl", "status", "smbd", "--no-pager"], check=False)


def configure_ssh(admin_user: str | None) -> None:
    sshd = Path("/etc/ssh/sshd_config")
    backup_file(sshd)
    dropin_dir = Path("/etc/ssh/sshd_config.d")
    dropin_dir.mkdir(parents=True, exist_ok=True)
    content = """# Hephaestus NAS baseline SSH hardening.
# Password auth is left enabled by default so first setup does not lock you out.
PermitRootLogin no
PubkeyAuthentication yes
X11Forwarding no
ClientAliveInterval 300
ClientAliveCountMax 2
"""
    (dropin_dir / "99-hephaestus-nas.conf").write_text(content, encoding="utf-8")
    if admin_user:
        run(["usermod", "-aG", "sudo,sambashare", admin_user], check=False)
    run(["sshd", "-t"])
    run(["systemctl", "restart", "ssh"])


def configure_samba(admin_user: str | None, create_default_shares: bool) -> None:
    smb_conf = Path("/etc/samba/smb.conf")
    backup_file(smb_conf)
    if create_default_shares:
        for path in [Path("/mnt/media_primary"), Path("/mnt/media_backup")]:
            path.mkdir(parents=True, exist_ok=True)
            run(["chgrp", "sambashare", str(path)], check=False)
            run(["chmod", "2775", str(path)], check=False)

        existing = smb_conf.read_text(encoding="utf-8", errors="ignore") if smb_conf.exists() else ""
        managed_pattern = re.compile(
            re.escape(SMB_MANAGED_BEGIN) + r".*?" + re.escape(SMB_MANAGED_END) + r"\n?",
            re.DOTALL,
        )
        if managed_pattern.search(existing):
            new = managed_pattern.sub(DEFAULT_SHARE_BLOCK, existing)
        else:
            new = existing.rstrip() + "\n\n" + DEFAULT_SHARE_BLOCK
        smb_conf.write_text(new, encoding="utf-8")
        log("Installed managed Samba shares for /mnt/media_primary and /mnt/media_backup")
    else:
        log("Skipped default Samba share creation. Use --create-default-shares to add media_primary/media_backup shares.")

    if admin_user:
        log(f"Adding {admin_user!r} to sambashare group. You still need to set a Samba password manually:")
        log(f"  sudo smbpasswd -a {admin_user}")
        run(["usermod", "-aG", "sambashare", admin_user], check=False)

    run(["testparm", "-s"], check=True)
    run(["systemctl", "restart", "smbd", "nmbd"], check=False)


def configure_firewall() -> None:
    # Safe LAN/NAS defaults. Does not expose anything to the public internet by itself.
    rules = [
        ["ufw", "allow", "OpenSSH"],
        ["ufw", "allow", "Samba"],
        ["ufw", "allow", "80/tcp", "comment", "CasaOS web UI"],
        ["ufw", "allow", "443/tcp", "comment", "HTTPS web UI if enabled"],
        ["ufw", "allow", "9090/tcp", "comment", "Cockpit web admin"],
    ]
    for rule in rules:
        run(rule, check=False)
    run(["ufw", "--force", "enable"], check=False)
    run(["ufw", "status", "verbose"], check=False)


def configure_unattended_upgrades() -> None:
    run(["dpkg-reconfigure", "-f", "noninteractive", "unattended-upgrades"], check=False)
    Path("/etc/apt/apt.conf.d/20auto-upgrades").write_text(
        'APT::Periodic::Update-Package-Lists "1";\n'
        'APT::Periodic::Unattended-Upgrade "1";\n'
        'APT::Periodic::AutocleanInterval "7";\n',
        encoding="utf-8",
    )
    log("Enabled unattended security upgrades.")


def install_tailscale(skip: bool) -> None:
    if skip:
        log("Skipped Tailscale install.")
        return
    if shutil.which("tailscale"):
        log("Tailscale already installed.")
    else:
        run(["bash", "-c", "curl -fsSL https://tailscale.com/install.sh | sh"])
    run(["systemctl", "enable", "--now", "tailscaled"], check=False)
    log("Tailscale installed. To authenticate this NAS, run:")
    log("  sudo tailscale up --ssh")
    log("If you do not want Tailscale SSH, use: sudo tailscale up")


def install_casaos(skip: bool) -> None:
    if skip:
        log("Skipped CasaOS install.")
        return
    if Path("/usr/bin/casaos").exists() or Path("/etc/systemd/system/casaos.service").exists():
        log("CasaOS appears to already be installed.")
        return
    run(["bash", "-c", "curl -fsSL https://get.casaos.io | bash"])
    run(["systemctl", "status", "casaos-gateway", "--no-pager"], check=False)


def print_disk_report() -> None:
    log("Current block devices. Verify media disks carefully before mounting/formatting anything:")
    run(["lsblk", "-o", "NAME,SIZE,TYPE,FSTYPE,LABEL,UUID,MOUNTPOINTS"], check=False)
    run(["df", "-hT"], check=False)


def final_report(hostname: str | None) -> None:
    host = hostname or platform.node()
    log("Bootstrap complete.")
    log(f"Try CasaOS from another machine: http://{host}.local/ or http://<NAS-LAN-IP>/")
    log(f"Try Cockpit admin panel: https://{host}.local:9090/ or https://<NAS-LAN-IP>:9090/")
    log("Next manual steps:")
    log("  1. Authenticate Tailscale: sudo tailscale up --ssh")
    log("  2. Set Samba password: sudo smbpasswd -a <your-user>")
    log("  3. Reconnect media drives if they were unplugged during OS install.")
    log("  4. Use lsblk/blkid to identify media drive UUIDs before editing /etc/fstab.")
    log("  5. Reboot once, then verify: systemctl status ssh smbd tailscaled casaos-gateway")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Bootstrap Debian headless into a CasaOS NAS.")
    parser.add_argument("--hostname", help="Set NAS hostname, e.g. dell-nas")
    parser.add_argument("--admin-user", help="Existing Debian user to add to sudo/sambashare groups")
    parser.add_argument("--create-default-shares", action="store_true", help="Create /mnt/media_primary and /mnt/media_backup Samba shares")
    parser.add_argument("--skip-casaos", action="store_true", help="Do not install CasaOS")
    parser.add_argument("--skip-tailscale", action="store_true", help="Do not install Tailscale")
    parser.add_argument("--no-firewall", action="store_true", help="Do not enable/configure ufw")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    require_root()
    check_debian_family()
    set_hostname(args.hostname)
    apt_install()
    configure_ssh(args.admin_user)
    configure_samba(args.admin_user, args.create_default_shares)
    enable_services()
    configure_unattended_upgrades()
    if not args.no_firewall:
        configure_firewall()
    install_tailscale(args.skip_tailscale)
    install_casaos(args.skip_casaos)
    print_disk_report()
    final_report(args.hostname)
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except KeyboardInterrupt:
        raise SystemExit("Interrupted.")
    except Exception as exc:
        log(f"ERROR: {exc}")
        raise