#!/usr/bin/env python3
"""
Alternative Debian Atomic Solution
Uses debootstrap/mmdebstrap + ostree instead of apt-ostree to create Debian Atomic systems

Features:
- Automatic detection of mmdebstrap (preferred) or debootstrap (fallback)
- mmdebstrap is faster and more reliable than debootstrap
- Automatic installation of mmdebstrap if not available
- Support for all major Debian desktop environments
- OSTree integration for atomic system management

Usage:
  python3 create-debian-atomic.py <variant> [work_dir] [options]
  
  Variants: minimal, gnome, plasma, cosmic, sway, budgie
  Options: --use-debootstrap, --help, -h
"""

import os
import sys
import subprocess
import tempfile
import shutil
import json
import yaml
from pathlib import Path
from typing import List, Dict, Any

class DebianAtomicBuilder:
    """Build Debian Atomic systems using debootstrap/mmdebstrap + ostree"""
    
    def __init__(self, work_dir: str = None, use_mmdebstrap: bool = True):
        self.work_dir = Path(work_dir) if work_dir else Path(tempfile.mkdtemp(prefix="debian-atomic-"))
        self.repo_dir = self.work_dir / "repo"
        self.rootfs_dir = self.work_dir / "rootfs"
        self.use_mmdebstrap = use_mmdebstrap
        
        # Check which bootstrap tool is available
        self.bootstrap_tool = self._detect_bootstrap_tool()
        
    def _detect_bootstrap_tool(self) -> str:
        """Detect which bootstrap tool is available and preferred"""
        if self.use_mmdebstrap and self._check_tool_available("mmdebstrap"):
            print("Using mmdebstrap (faster and more reliable)")
            return "mmdebstrap"
        elif self._check_tool_available("debootstrap"):
            print("Using debootstrap (fallback option)")
            return "debootstrap"
        else:
            print("Warning: No bootstrap tool found, attempting to install mmdebstrap")
            if self._install_mmdebstrap():
                return "mmdebstrap"
            else:
                raise RuntimeError("No bootstrap tool available and could not install mmdebstrap")
    
    def _check_tool_available(self, tool: str) -> bool:
        """Check if a tool is available in PATH"""
        try:
            subprocess.run([tool, "--version"], capture_output=True, check=True)
            return True
        except (subprocess.CalledProcessError, FileNotFoundError):
            return False
    
    def _install_mmdebstrap(self) -> bool:
        """Attempt to install mmdebstrap if not available"""
        try:
            print("Attempting to install mmdebstrap...")
            # Try different package managers
            install_cmds = [
                ["apt-get", "update", "-y"],
                ["apt-get", "install", "-y", "mmdebstrap"]
            ]
            
            for cmd in install_cmds:
                result = subprocess.run(cmd, capture_output=True, text=True, check=True)
                print(f"Successfully ran: {' '.join(cmd)}")
            
            return self._check_tool_available("mmdebstrap")
        except subprocess.CalledProcessError as e:
            print(f"Failed to install mmdebstrap: {e}")
            return False
        
    def create_ostree_repo(self) -> bool:
        """Create and initialize OSTree repository"""
        try:
            print(f"Creating OSTree repository at {self.repo_dir}")
            self.repo_dir.mkdir(parents=True, exist_ok=True)
            
            # Initialize OSTree repository
            result = subprocess.run([
                "ostree", "init", "--repo", str(self.repo_dir), "--mode=bare"
            ], capture_output=True, text=True, check=True)
            
            print("OSTree repository initialized successfully")
            return True
            
        except subprocess.CalledProcessError as e:
            print(f"Failed to initialize OSTree repository: {e}")
            print(f"stdout: {e.stdout}")
            print(f"stderr: {e.stderr}")
            return False
        except Exception as e:
            print(f"Error creating OSTree repository: {e}")
            return False
    
    def create_rootfs(self, variant: str = "minimal") -> bool:
        """Create rootfs using detected bootstrap tool"""
        try:
            print(f"Creating rootfs for variant: {variant} using {self.bootstrap_tool}")
            self.rootfs_dir.mkdir(parents=True, exist_ok=True)
            
            # Get package list based on variant
            packages = self.get_packages_for_variant(variant)
            
            if self.bootstrap_tool == "mmdebstrap":
                # Use mmdebstrap (faster and more reliable)
                cmd = [
                    "mmdebstrap", 
                    "--variant=minbase", 
                    "--include=" + ",".join(packages),
                    "--aptopt=Acquire::Check-Valid-Until=false",  # Handle expired Release files
                    "--aptopt=Acquire::Check-Date=false",         # Disable date checking
                    "trixie", 
                    str(self.rootfs_dir), 
                    "http://deb.debian.org/debian"
                ]
            else:
                # Use debootstrap (fallback)
                cmd = [
                    "debootstrap", 
                    "--variant=minbase", 
                    "--include=" + ",".join(packages),
                    "trixie", 
                    str(self.rootfs_dir), 
                    "http://deb.debian.org/debian"
                ]
            
            print(f"Running: {' '.join(cmd)}")
            result = subprocess.run(cmd, capture_output=True, text=True, check=True)
            
            print("Rootfs created successfully")
            return True
            
        except subprocess.CalledProcessError as e:
            print(f"Failed to create rootfs: {e}")
            print(f"stdout: {e.stdout}")
            print(f"stderr: {e.stderr}")
            return False
        except Exception as e:
            print(f"Error creating rootfs: {e}")
            return False
    
    def get_packages_for_variant(self, variant: str) -> List[str]:
        """Get package list for a specific variant"""
        # Start with minimal set of packages that definitely exist
        base_packages = [
            "systemd", "dbus", "sudo", "bash",
            "coreutils", "util-linux", "procps"
        ]
        
        if variant == "minimal":
            return base_packages + [
                "less", "vim-tiny", "wget", "curl", "ca-certificates",
                "gnupg", "iproute2", "net-tools", "openssh-client",
                "openssh-server", "htop", "rsync", "tar", "gzip", "unzip"
            ]
        elif variant == "gnome":
            return base_packages + [
                "gnome-shell", "gnome-session", "gnome-terminal",
                "gnome-control-center", "gnome-settings-daemon",
                "gnome-backgrounds", "gnome-themes-extra",
                "adwaita-icon-theme", "gdm3", "gnome-initial-setup",
                "nautilus", "gnome-software", "gnome-tweaks"
            ]
        elif variant == "plasma":
            return base_packages + [
                "plasma-desktop", "plasma-workspace", "plasma-nm",
                "plasma-pa", "kde-config", "kde-runtime", "kde-standard",
                "sddm", "kwin-x11", "dolphin", "konsole", "kate",
                "kdeconnect", "plasma-browser-integration"
            ]
        elif variant == "cosmic":
            return base_packages + [
                "pop-desktop", "pop-shell", "gnome-shell",
                "gnome-session", "gnome-terminal", "gnome-control-center",
                "gnome-settings-daemon", "adwaita-icon-theme", "gdm3",
                "nautilus", "gnome-software", "pop-gtk-theme"
            ]
        elif variant == "sway":
            return base_packages + [
                "sway", "swaybg", "swayidle", "swaylock",
                "waybar", "wofi", "foot", "grim", "slurp",
                "wl-clipboard", "mako", "swaymsg", "sway-input"
            ]
        elif variant == "budgie":
            return base_packages + [
                "budgie-desktop", "budgie-desktop-view",
                "budgie-panel", "budgie-menu", "budgie-run-dialog",
                "budgie-screenshot", "budgie-session",
                "gnome-session", "gdm3", "adwaita-icon-theme"
            ]
        else:
            return base_packages
    
    def create_ostree_commit(self, variant: str, ref: str) -> str:
        """Create OSTree commit from rootfs"""
        try:
            print(f"Creating OSTree commit for variant: {variant}")
            
            # Create commit
            cmd = [
                "ostree", "commit", "--repo", str(self.repo_dir),
                "--branch", ref, "--tree", f"dir={self.rootfs_dir}",
                f"--subject=Debian Atomic {variant} variant",
                f"--body=Debian Trixie Atomic system with {variant} desktop"
            ]
            
            print(f"Running: {' '.join(cmd)}")
            result = subprocess.run(cmd, capture_output=True, text=True, check=True)
            
            # Extract commit hash from output
            commit_hash = result.stdout.strip()
            print(f"OSTree commit created: {commit_hash}")
            return commit_hash
            
        except subprocess.CalledProcessError as e:
            print(f"Failed to create OSTree commit: {e}")
            print(f"stdout: {e.stdout}")
            print(f"stderr: {e.stderr}")
            return None
        except Exception as e:
            print(f"Error creating OSTree commit: {e}")
            return None
    
    def post_process_rootfs(self, variant: str) -> bool:
        """Post-process rootfs after debootstrap"""
        try:
            print(f"Post-processing rootfs for {variant} variant")
            
            # Create essential directories
            essential_dirs = [
                "/etc/apt-ostree",
                "/var/lib/apt-ostree", 
                "/usr/lib/bootc",
                "/root/.ssh",
                "/etc/systemd/system",
                "/etc/systemd/user"
            ]
            
            for dir_path in essential_dirs:
                full_path = self.rootfs_dir / dir_path.lstrip("/")
                full_path.mkdir(parents=True, exist_ok=True)
            
            # Set up basic systemd services
            if variant in ["gnome", "plasma", "cosmic", "budgie"]:
                # Enable display manager
                if variant == "plasma":
                    sddm_service = self.rootfs_dir / "etc/systemd/system/display-manager.service"
                    sddm_service.write_text("[Unit]\nDescription=SDDM Display Manager\n\n[Service]\nExecStart=/usr/bin/sddm\nRestart=always\n\n[Install]\nWantedBy=graphical.target\n")
                elif variant in ["gnome", "cosmic", "budgie"]:
                    gdm_service = self.rootfs_dir / "etc/systemd/system/display-manager.service"
                    gdm_service.write_text("[Unit]\nDescription=GDM Display Manager\n\n[Service]\nExecStart=/usr/sbin/gdm3\nRestart=always\n\n[Install]\nWantedBy=graphical.target\n")
            
            # Enable SSH service
            ssh_service = self.rootfs_dir / "etc/systemd/system/sshd.service"
            ssh_service.write_text("[Unit]\nDescription=OpenSSH Server\n\n[Service]\nExecStart=/usr/sbin/sshd -D\nRestart=always\n\n[Install]\nWantedBy=multi-user.target\n")
            
            print("Rootfs post-processing completed")
            return True
            
        except Exception as e:
            print(f"Error during post-processing: {e}")
            return False
    
    def build_variant(self, variant: str) -> bool:
        """Build a complete Debian Atomic variant"""
        try:
            print(f"Building Debian Atomic variant: {variant}")
            
            # Create OSTree repository
            if not self.create_ostree_repo():
                return False
            
            # Create rootfs
            if not self.create_rootfs(variant):
                return False
            
            # Post-process rootfs
            if not self.post_process_rootfs(variant):
                return False
            
            # Create OSTree commit
            ref = f"debian/14/x86_64/{variant}"
            commit_hash = self.create_ostree_commit(variant, ref)
            
            if not commit_hash:
                return False
            
            print(f"Successfully built {variant} variant")
            print(f"Repository: {self.repo_dir}")
            print(f"Reference: {ref}")
            print(f"Commit: {commit_hash}")
            
            return True
            
        except Exception as e:
            print(f"Error building variant {variant}: {e}")
            return False
    
    def list_variants(self) -> List[str]:
        """List all supported variants"""
        return ["minimal", "gnome", "plasma", "cosmic", "sway", "budgie"]
    
    def get_bootstrap_info(self) -> Dict[str, Any]:
        """Get information about the bootstrap tool being used"""
        return {
            "tool": self.bootstrap_tool,
            "preferred": self.use_mmdebstrap,
            "available": self._check_tool_available(self.bootstrap_tool),
            "version": self._get_tool_version()
        }
    
    def _get_tool_version(self) -> str:
        """Get version of the bootstrap tool"""
        try:
            result = subprocess.run([self.bootstrap_tool, "--version"], 
                                  capture_output=True, text=True, check=True)
            return result.stdout.strip().split('\n')[0]
        except:
            return "Unknown"
    
    def cleanup(self):
        """Clean up temporary files"""
        try:
            if self.work_dir.exists():
                shutil.rmtree(self.work_dir)
                print(f"Cleaned up {self.work_dir}")
        except Exception as e:
            print(f"Error during cleanup: {e}")

def main():
    """Main function"""
    if len(sys.argv) < 2 or "--help" in sys.argv or "-h" in sys.argv:
        print("Alternative Debian Atomic Builder")
        print("=================================")
        print("Usage: python3 create-debian-atomic.py <variant> [work_dir] [options]")
        print("\nVariants:")
        for variant in DebianAtomicBuilder().list_variants():
            print(f"  {variant}")
        print("\nOptions:")
        print("  --use-debootstrap    Force use of debootstrap instead of mmdebstrap")
        print("  --help, -h           Show this help message")
        print("\nExamples:")
        print("  python3 create-debian-atomic.py minimal")
        print("  python3 create-debian-atomic.py gnome /tmp/my-work-dir")
        print("  python3 create-debian-atomic.py plasma --use-debootstrap")
        sys.exit(0)
    
    variant = sys.argv[1]
    work_dir = None
    use_mmdebstrap = True
    
    # Parse arguments
    for i, arg in enumerate(sys.argv[2:], 2):
        if arg == "--use-debootstrap":
            use_mmdebstrap = False
        elif arg in ["--help", "-h"]:
            continue  # Already handled above
        elif not arg.startswith("--"):
            work_dir = arg
            break
    
    builder = DebianAtomicBuilder(work_dir=work_dir, use_mmdebstrap=use_mmdebstrap)
    supported_variants = builder.list_variants()
    
    if variant not in supported_variants:
        print(f"Unknown variant: {variant}")
        print(f"Supported variants: {', '.join(supported_variants)}")
        sys.exit(1)
    
    try:
        # Show bootstrap tool information
        bootstrap_info = builder.get_bootstrap_info()
        print(f"\n🔧 Bootstrap Tool Information:")
        print(f"  Tool: {bootstrap_info['tool']}")
        print(f"  Version: {bootstrap_info['version']}")
        print(f"  Preferred: {'mmdebstrap' if bootstrap_info['preferred'] else 'debootstrap'}")
        
        success = builder.build_variant(variant)
        if success:
            print(f"\n✅ Successfully built Debian Atomic {variant} variant")
            print(f"Repository location: {builder.repo_dir}")
            print(f"To deploy: ostree admin deploy --os {variant} {builder.repo_dir}")
            print(f"To list refs: ostree refs --repo={builder.repo_dir}")
        else:
            print(f"\n❌ Failed to build Debian Atomic {variant} variant")
            sys.exit(1)
    finally:
        if not work_dir:  # Only cleanup if we created a temp directory
            builder.cleanup()

if __name__ == "__main__":
    main()