#!/usr/bin/env python3
"""
TotiOS Installer — Flash Tool
Lädt das neueste TotiOS Image herunter und schreibt es direkt auf USB-Stick oder SD-Karte.
Cross-platform: macOS, Linux, Windows.

Usage:
  python3 totios-installer.py                    # Interactive mode
  python3 totios-installer.py --list             # List available images
  python3 totios-installer.py --variant headless --arch aarch64 --device /dev/sdX
"""

import argparse
import hashlib
import json
import os
import platform
import shutil
import subprocess
import sys
import tempfile
import urllib.request
import urllib.error

# ===== CONFIG =====
BASE_URL = "https://mercury-og.tailaa8f66.ts.net/totios/downloads"
MANIFEST_URL = f"{BASE_URL}/manifest.json"

# ===== COLORS =====
class C:
    R = "\033[0m"
    BOLD = "\033[1m"
    BLUE = "\033[34m"
    GREEN = "\033[32m"
    YELLOW = "\033[33m"
    RED = "\033[31m"
    CYAN = "\033[36m"
    DIM = "\033[2m"

def info(msg):
    print(f"{C.BLUE}[i]{C.R} {msg}")

def ok(msg):
    print(f"{C.GREEN}[+]{C.R} {msg}")

def warn(msg):
    print(f"{C.YELLOW}[!]{C.R} {msg}")

def err(msg):
    print(f"{C.RED}[x]{C.R} {msg}")

def header(msg):
    print(f"\n{C.BOLD}{C.CYAN}{'='*50}{C.R}")
    print(f"{C.BOLD}{C.CYAN} {msg}{C.R}")
    print(f"{C.BOLD}{C.CYAN}{'='*50}{C.R}\n")

# ===== DOWNLOAD =====
def download_progress(url, dest):
    """Download with progress bar."""
    req = urllib.request.Request(url, headers={"User-Agent": "TotiOS-Installer/1.0"})
    try:
        resp = urllib.request.urlopen(req)
    except urllib.error.URLError as e:
        err(f"Download failed: {e}")
        return False

    total = int(resp.headers.get("Content-Length", 0))
    downloaded = 0
    chunk_size = 1024 * 1024  # 1MB

    with open(dest, "wb") as f:
        while True:
            chunk = resp.read(chunk_size)
            if not chunk:
                break
            f.write(chunk)
            downloaded += len(chunk)
            if total > 0:
                pct = downloaded * 100 // total
                bar = "=" * (pct // 2) + " " * (50 - pct // 2)
                mb_down = downloaded // (1024 * 1024)
                mb_total = total // (1024 * 1024)
                print(f"\r  {C.DIM}[{bar}]{C.R} {pct:3d}%  {mb_down}/{mb_total} MB", end="", flush=True)
            else:
                mb_down = downloaded // (1024 * 1024)
                print(f"\r  {mb_down} MB downloaded...", end="", flush=True)

    print()
    return True

def verify_sha256(filepath, expected_hash):
    """Verify SHA256 checksum of a file."""
    h = hashlib.sha256()
    with open(filepath, "rb") as f:
        while True:
            chunk = f.read(1024 * 1024)
            if not chunk:
                break
            h.update(chunk)
    actual = h.hexdigest()
    return actual.lower() == expected_hash.lower()

# ===== DEVICE DETECTION =====
def list_devices():
    """List removable storage devices (USB sticks, SD cards)."""
    system = platform.system()

    if system == "Darwin":  # macOS
        result = subprocess.run(
            ["diskutil", "list", "external"],
            capture_output=True, text=True
        )
        devices = []
        current_disk = None
        for line in result.stdout.split("\n"):
            if line.startswith("/dev/disk"):
                current_disk = line.split()[0].replace("/dev/", "")
                devices.append({"id": line.split()[0], "name": "External Disk", "size": "?"})
            elif current_disk and "NAME" not in line and line.strip():
                parts = line.split()
                if len(parts) >= 2:
                    name = " ".join(parts[:-1]) if len(parts) > 2 else parts[0]
                    devices[-1]["name"] = name
        return devices

    elif system == "Linux":
        devices = []
        # Check /dev for sd* and mmcblk*
        for dev in sorted(os.listdir("/dev")):
            if (dev.startswith("sd") and dev[-1].isdigit() == False and len(dev) <= 4) or \
               dev.startswith("mmcblk") and "boot" not in dev and "p" not in dev[7:]:
                path = f"/dev/{dev}"
                # Check if removable
                try:
                    with open(f"/sys/block/{dev}/removable") as f:
                        removable = f.read().strip() == "1"
                    if removable:
                        # Get size
                        with open(f"/sys/block/{dev}/size") as f:
                            sectors = int(f.read().strip())
                        size_gb = sectors * 512 / (1024**3)
                        # Get model/name
                        try:
                            with open(f"/sys/block/{dev}/device/model") as f:
                                name = f.read().strip()
                        except:
                            name = dev
                        devices.append({"id": path, "name": name, "size": f"{size_gb:.1f} GB"})
                except:
                    pass
        return devices

    elif system == "Windows":
        result = subprocess.run(
            ["wmic", "diskdrive", "get", "DeviceID,Model,Size,InterfaceType"],
            capture_output=True, text=True
        )
        # Parse WMIC output
        devices = []
        for line in result.stdout.strip().split("\n")[1:]:
            parts = line.strip().split()
            if len(parts) >= 2:
                devices.append({"id": parts[-1], "name": " ".join(parts[:-1]), "size": "?"})
        return devices

    return []

def unmount_device(device_path):
    """Unmount all partitions on a device."""
    system = platform.system()

    if system == "Darwin":
        subprocess.run(["diskutil", "unmountDisk", device_path], capture_output=True)
    elif system == "Linux":
        # Unmount all partitions
        result = subprocess.run(["lsblk", "-no", "MOUNTPOINT", device_path],
                               capture_output=True, text=True)
        for mountpoint in result.stdout.strip().split("\n"):
            mountpoint = mountpoint.strip()
            if mountpoint:
                subprocess.run(["umount", mountpoint], capture_output=True)
    elif system == "Windows":
        # Windows doesn't need unmount before dd
        pass

# ===== FLASH =====
def flash_image(image_path, device_path):
    """Write image to device using dd (macOS/Linux) or dd for Windows."""
    system = platform.system()

    if system == "Darwin":
        # macOS: use /dev/rdisk* for faster writes
        rdisk_path = device_path.replace("/dev/disk", "/dev/rdisk")
        unmount_device(device_path)

        info(f"Flashing to {rdisk_path}...")
        cmd = ["dd", f"if={image_path}", f"of={rdisk_path}", "bs=4m"]
        result = subprocess.run(cmd, capture_output=True, text=True)
        if result.returncode != 0:
            err(f"dd failed: {result.stderr}")
            return False

        # Decompress if .gz
        if image_path.endswith(".gz"):
            info("Decompressing on-the-fly...")
            cmd = ["sh", "-c", f"gunzip -c '{image_path}' | dd of='{rdisk_path}' bs=4m"]
            result = subprocess.run(cmd, capture_output=True, text=True)
            if result.returncode != 0:
                err(f"gunzip+dd failed: {result.stderr}")
                return False

        ok("Flash complete!")
        subprocess.run(["diskutil", "eject", device_path], capture_output=True)
        return True

    elif system == "Linux":
        unmount_device(device_path)

        if image_path.endswith(".gz"):
            info(f"Decompressing and flashing to {device_path}...")
            cmd = ["sh", "-c", f"gunzip -c '{image_path}' | dd of='{device_path}' bs=4M status=progress"]
        else:
            info(f"Flashing to {device_path}...")
            cmd = ["dd", f"if={image_path}", f"of={device_path}", "bs=4M", "status=progress"]

        result = subprocess.run(cmd, capture_output=True, text=True)
        if result.returncode != 0:
            err(f"dd failed: {result.stderr}")
            return False

        ok("Flash complete!")
        subprocess.run(["sync"])
        return True

    elif system == "Windows":
        # Use dd for Windows or a Python-based approach
        warn("Windows flash: using Python file copy (slower but works)")
        unmount_device(device_path)
        try:
            if image_path.endswith(".gz"):
                import gzip
                with gzip.open(image_path, "rb") as src, open(device_path, "wb") as dst:
                    shutil.copyfileobj(src, dst, length=1024*1024)
            else:
                with open(image_path, "rb") as src, open(device_path, "wb") as dst:
                    shutil.copyfileobj(src, dst, length=1024*1024)
            ok("Flash complete!")
            return True
        except Exception as e:
            err(f"Flash failed: {e}")
            return False

    return False

# ===== MANIFEST =====
def fetch_manifest():
    """Download manifest from the TotiOS download server."""
    req = urllib.request.Request(MANIFEST_URL, headers={"User-Agent": "TotiOS-Installer/1.0"})
    try:
        resp = urllib.request.urlopen(req, timeout=15)
        return json.loads(resp.read())
    except Exception as e:
        err(f"Cannot fetch manifest: {e}")
        return None

# ===== INTERACTIVE =====
def interactive_mode():
    """Interactive wizard: select variant → select device → flash."""
    header("TotiOS Installer")

    # 1. Fetch manifest
    info("Fetching latest image catalog...")
    manifest = fetch_manifest()
    if not manifest:
        err("Cannot reach TotiOS download server. Are you online?")
        return 1

    images = manifest.get("images", [])
    ok(f"Found {len(images)} images (v{manifest.get('version', '?')})")

    # 2. Select variant
    variants = []
    for img in images:
        v = img["variant"]
        if v not in variants:
            variants.append(v)

    print(f"\n{C.BOLD}Available variants:{C.R}")
    for i, v in enumerate(variants, 1):
        descs = set(img["description"] for img in images if img["variant"] == v)
        print(f"  {C.CYAN}{i}{C.R}. {C.BOLD}{v}{C.R} — {list(descs)[0]}")

    while True:
        try:
            choice = int(input(f"\nSelect variant [1-{len(variants)}]: "))
            if 1 <= choice <= len(variants):
                selected_variant = variants[choice - 1]
                break
        except (ValueError, KeyboardInterrupt):
            err("Aborted.")
            return 1

    # 3. Select architecture
    archs = [img["arch"] for img in images if img["variant"] == selected_variant]
    print(f"\n{C.BOLD}Available architectures for {selected_variant}:{C.R}")
    for i, a in enumerate(archs, 1):
        device = next(img["device"] for img in images if img["variant"] == selected_variant and img["arch"] == a)
        print(f"  {C.CYAN}{i}{C.R}. {C.BOLD}{a}{C.R} — {device}")

    while True:
        try:
            choice = int(input(f"\nSelect architecture [1-{len(archs)}]: "))
            if 1 <= choice <= len(archs):
                selected_arch = archs[choice - 1]
                break
        except (ValueError, KeyboardInterrupt):
            err("Aborted.")
            return 1

    # 4. Find matching image
    selected = next((img for img in images if img["variant"] == selected_variant and img["arch"] == selected_arch), None)
    if not selected:
        err("Image not found.")
        return 1

    filename = selected["filename"]
    url = f"{BASE_URL}/{filename}"
    is_gz = filename.endswith(".gz")

    ok(f"Selected: {filename}")
    info(f"Description: {selected['description']}")

    # 5. List removable devices
    print(f"\n{C.BOLD}Detecting removable storage devices...{C.R}")
    devices = list_devices()

    if not devices:
        warn("No removable devices detected automatically.")
        print("  You can specify the device path manually.")
        manual = input("Enter device path (e.g. /dev/sdb, /dev/disk2): ").strip()
        if not manual:
            err("No device selected.")
            return 1
        selected_device = manual
    else:
        print(f"\n{C.BOLD}Found {len(devices)} device(s):{C.R}")
        for i, d in enumerate(devices, 1):
            print(f"  {C.CYAN}{i}{C.R}. {d['id']} — {d['name']} ({d.get('size', '?')})")
        print(f"  {C.CYAN}m{C.R}. Enter device path manually")

        while True:
            try:
                choice = input(f"\nSelect device [1-{len(devices)}/m]: ").strip()
                if choice.lower() == "m":
                    selected_device = input("Enter device path: ").strip()
                    if not selected_device:
                        err("No device selected.")
                        return 1
                    break
                choice = int(choice)
                if 1 <= choice <= len(devices):
                    selected_device = devices[choice - 1]["id"]
                    break
            except (ValueError, KeyboardInterrupt):
                err("Aborted.")
                return 1

    # 6. Confirm
    header("CONFIRMATION")
    print(f"  Image:   {filename}")
    print(f"  Source:  {url}")
    print(f"  Target:  {selected_device}")
    print(f"  Type:    {'SD Card Image (.img.gz)' if is_gz else 'ISO (bootable)'}")
    print()
    warn("ALL DATA ON THE TARGET DEVICE WILL BE DESTROYED!")
    confirm = input(f"\nType 'YES' to proceed: ").strip()
    if confirm != "YES":
        err("Aborted by user.")
        return 1

    # 7. Download
    header("DOWNLOAD")
    tmpdir = tempfile.mkdtemp(prefix="totios_")
    dest = os.path.join(tmpdir, filename)
    info(f"Downloading {filename}...")
    if not download_progress(url, dest):
        err("Download failed.")
        shutil.rmtree(tmpdir, ignore_errors=True)
        return 1
    ok("Download complete!")

    # 8. Verify checksum
    info("Verifying SHA256 checksum...")
    sha_url = f"{BASE_URL}/{filename}.sha256"
    try:
        req = urllib.request.Request(sha_url, headers={"User-Agent": "TotiOS-Installer/1.0"})
        resp = urllib.request.urlopen(req, timeout=15)
        expected_hash = resp.read().decode().strip().split()[0]
    except:
        expected_hash = None
        warn("No checksum available — skipping verification.")

    if expected_hash:
        if verify_sha256(dest, expected_hash):
            ok("Checksum verified!")
        else:
            err("Checksum mismatch! The downloaded file may be corrupted.")
            shutil.rmtree(tmpdir, ignore_errors=True)
            return 1

    # 9. Flash
    header("FLASH")
    if flash_image(dest, selected_device):
        ok("TotiOS has been flashed successfully!")
        print(f"\n{C.BOLD}Next steps:{C.R}")
        if is_gz:
            print("  1. Insert the SD card into your Raspberry Pi")
            print("  2. Connect the Pi to power and network")
            print("  3. TotiOS will auto-install on first boot (~5 min)")
            print("  4. After reboot, connect via Tailscale or Mercury Remote")
        else:
            print("  1. Insert the USB stick / burn ISO to USB")
            print("  2. Boot from the USB stick")
            print("  3. TotiOS will auto-install on first boot")
        print(f"\n{C.GREEN}Done!{C.R}")
        shutil.rmtree(tmpdir, ignore_errors=True)
        return 0
    else:
        err("Flash failed!")
        shutil.rmtree(tmpdir, ignore_errors=True)
        return 1

# ===== MAIN =====
def main():
    parser = argparse.ArgumentParser(description="TotiOS Installer — Flash tool for USB/SD cards")
    parser.add_argument("--list", action="store_true", help="List available images")
    parser.add_argument("--variant", choices=["headless", "console", "gnome", "ui"], help="Select variant")
    parser.add_argument("--arch", choices=["aarch64", "armv7l", "x86_64"], help="Select architecture")
    parser.add_argument("--device", help="Target device path (e.g. /dev/sdb)")
    parser.add_argument("--yes", action="store_true", help="Skip confirmation")
    args = parser.parse_args()

    if args.list:
        header("Available TotiOS Images")
        manifest = fetch_manifest()
        if not manifest:
            err("Cannot reach download server.")
            return 1
        images = manifest.get("images", [])
        print(f"Version: {manifest.get('version', '?')}")
        print(f"Updated: {manifest.get('updated', '?')}")
        print()
        for img in images:
            print(f"  {C.BOLD}{img['variant']:<8}{C.R} {img['arch']:<8} {img['device']:<25} {img['filename']}")
        return 0

    if args.variant and args.arch and args.device:
        # Non-interactive mode
        manifest = fetch_manifest()
        if not manifest:
            return 1
        images = manifest.get("images", [])
        selected = next((img for img in images if img["variant"] == args.variant and img["arch"] == args.arch), None)
        if not selected:
            err(f"No image for {args.variant}/{args.arch}")
            return 1

        filename = selected["filename"]
        url = f"{BASE_URL}/{filename}"

        if not args.yes:
            header("CONFIRMATION")
            print(f"  Image:   {filename}")
            print(f"  Target:  {args.device}")
            warn("ALL DATA ON THE TARGET DEVICE WILL BE DESTROYED!")
            confirm = input("Type 'YES' to proceed: ").strip()
            if confirm != "YES":
                err("Aborted.")
                return 1

        tmpdir = tempfile.mkdtemp(prefix="totios_")
        dest = os.path.join(tmpdir, filename)
        info(f"Downloading {filename}...")
        if not download_progress(url, dest):
            return 1
        ok("Download complete!")

        if flash_image(dest, args.device):
            ok("Flash complete!")
            return 0
        else:
            return 1

    # Default: interactive mode
    return interactive_mode()

if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        print(f"\n{C.RED}Aborted.{C.R}")
        sys.exit(1)