#!/usr/bin/env python3
"""
Debian Filesystem Stage for osbuild

This stage handles Debian filesystem setup, OSTree integration points, and permission
handling. It creates the filesystem structure needed for immutable Debian systems.

Author: Debian bootc-image-builder team
License: Same as original bootc-image-builder
"""

import os
import subprocess
import tempfile
import json
import glob
import shutil
import pwd
import grp
from typing import Dict, List, Optional, Any
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class DebianFilesystemStage:
    """
    osbuild stage for Debian filesystem setup and OSTree integration.
    
    This stage handles:
    - Debian filesystem layout setup
    - OSTree integration points
    - Permission and ownership handling
    - /home -> /var/home symlink creation
    """
    
    def __init__(self, options: Dict[str, Any]):
        """
        Initialize the Debian Filesystem stage with configuration options.
        
        Args:
            options: Dictionary containing stage configuration
                - rootfs_type: Root filesystem type (e.g., 'ext4')
                - ostree_integration: Whether to enable OSTree integration
                - home_symlink: Whether to create /home -> /var/home symlink
                - users: List of users to create
                - groups: List of groups to create
                - permissions: Custom permission mappings
        """
        self.options = options
        self.rootfs_type = options.get('rootfs_type', 'ext4')
        self.ostree_integration = options.get('ostree_integration', True)
        self.home_symlink = options.get('home_symlink', True)
        self.users = options.get('users', [])
        self.groups = options.get('groups', [])
        self.permissions = options.get('permissions', {})
        
        logger.info(f"Debian Filesystem Stage initialized")
        logger.info(f"Root filesystem type: {self.rootfs_type}")
        logger.info(f"OSTree integration: {self.ostree_integration}")
        logger.info(f"Home symlink: {self.home_symlink}")
    
    def run(self, context) -> None:
        """
        Execute the Debian Filesystem stage within the osbuild context.
        
        Args:
            context: osbuild context providing chroot access
        """
        logger.info("Starting Debian Filesystem stage execution")
        
        try:
            # Step 1: Set up basic filesystem structure
            self._setup_filesystem_structure(context)
            
            # Step 2: Create OSTree integration points
            if self.ostree_integration:
                self._setup_ostree_integration(context)
            
            # Step 3: Set up /home -> /var/home symlink
            if self.home_symlink:
                self._setup_home_symlink(context)
            
            # Step 4: Create users and groups
            self._setup_users_and_groups(context)
            
            # Step 5: Set up permissions
            self._setup_permissions(context)
            
            # Step 6: Configure system directories
            self._configure_system_directories(context)
            
            logger.info("Debian Filesystem stage completed successfully")
            
        except Exception as e:
            logger.error(f"Debian Filesystem stage failed: {e}")
            raise
    
    def _setup_filesystem_structure(self, context) -> None:
        """
        Set up basic Debian filesystem structure.
        
        Args:
            context: osbuild context
        """
        logger.info("Setting up Debian filesystem structure")
        
        # Create essential directories
        essential_dirs = [
            "bin", "boot", "dev", "etc", "home", "lib", "lib64", "media", "mnt",
            "opt", "proc", "root", "run", "sbin", "srv", "sys", "tmp", "usr", "var"
        ]
        
        for directory in essential_dirs:
            dir_path = os.path.join(context.root, directory)
            os.makedirs(dir_path, exist_ok=True)
            logger.debug(f"Created directory: {dir_path}")
        
        # Create /usr subdirectories
        usr_dirs = [
            "bin", "lib", "lib64", "sbin", "share", "src", "local"
        ]
        
        for directory in usr_dirs:
            dir_path = os.path.join(context.root, "usr", directory)
            os.makedirs(dir_path, exist_ok=True)
            logger.debug(f"Created /usr directory: {dir_path}")
        
        # Create /var subdirectories
        var_dirs = [
            "cache", "lib", "local", "lock", "log", "opt", "run", "spool", "tmp"
        ]
        
        for directory in var_dirs:
            dir_path = os.path.join(context.root, "var", directory)
            os.makedirs(dir_path, exist_ok=True)
            logger.debug(f"Created /var directory: {dir_path}")
        
        # Create /etc subdirectories
        etc_dirs = [
            "apt", "default", "init.d", "network", "systemd", "udev"
        ]
        
        for directory in etc_dirs:
            dir_path = os.path.join(context.root, "etc", directory)
            os.makedirs(dir_path, exist_ok=True)
            logger.debug(f"Created /etc directory: {dir_path}")
        
        logger.info("Basic filesystem structure created")
    
    def _setup_ostree_integration(self, context) -> None:
        """
        Set up OSTree integration points in the filesystem.
        
        Args:
            context: osbuild context
        """
        logger.info("Setting up OSTree integration points")
        
        # Create OSTree directories
        ostree_dirs = [
            "ostree",
            "usr/lib/ostree-boot",
            "usr/lib/ostree-boot/scripts",
            "etc/ostree",
            "etc/ostree/remotes.d"
        ]
        
        for directory in ostree_dirs:
            dir_path = os.path.join(context.root, directory)
            os.makedirs(dir_path, exist_ok=True)
            logger.debug(f"Created OSTree directory: {dir_path}")
        
        # Create OSTree configuration
        ostree_conf = os.path.join(context.root, "etc", "ostree", "ostree.conf")
        ostree_content = [
            "[core]",
            "repo_mode=bare-user",
            "",
            "[sysroot]",
            "readonly=true",
            "bootloader=grub2",
            ""
        ]
        
        with open(ostree_conf, 'w') as f:
            f.write('\n'.join(ostree_content))
        
        os.chmod(ostree_conf, 0o644)
        logger.info(f"Created OSTree configuration: {ostree_conf}")
        
        # Create OSTree remote configuration
        remote_conf = os.path.join(context.root, "etc", "ostree", "remotes.d", "ostree.conf")
        remote_content = [
            "[remote \"ostree\"]",
            "url=https://ostree.example.com/repo",
            "gpg-verify=false",
            ""
        ]
        
        with open(remote_conf, 'w') as f:
            f.write('\n'.join(remote_content))
        
        os.chmod(remote_conf, 0o644)
        logger.info(f"Created OSTree remote configuration: {remote_conf}")
    
    def _setup_home_symlink(self, context) -> None:
        """
        Set up /home -> /var/home symlink for immutable architecture.
        
        Args:
            context: osbuild context
        """
        logger.info("Setting up /home -> /var/home symlink")
        
        # Create /var/home directory
        var_home = os.path.join(context.root, "var", "home")
        os.makedirs(var_home, exist_ok=True)
        
        # Set proper permissions for /var/home
        os.chmod(var_home, 0o755)
        
        # Remove existing /home if it exists (directory, file, or symlink)
        home_path = os.path.join(context.root, "home")
        if os.path.exists(home_path):
            if os.path.isdir(home_path):
                shutil.rmtree(home_path)
            else:
                os.remove(home_path)
        
        # Create symlink from /home to /var/home
        try:
            os.symlink("../var/home", home_path)
            logger.info(f"Created symlink: {home_path} -> ../var/home")
        except OSError as e:
            logger.error(f"Failed to create symlink: {e}")
            raise
    
    def _setup_users_and_groups(self, context) -> None:
        """
        Set up users and groups in the filesystem.
        
        Args:
            context: osbuild context
        """
        logger.info("Setting up users and groups")
        
        # Create default groups
        default_groups = [
            ("root", 0),
            ("daemon", 1),
            ("bin", 2),
            ("sys", 3),
            ("adm", 4),
            ("tty", 5),
            ("disk", 6),
            ("lp", 7),
            ("mail", 8),
            ("news", 9),
            ("uucp", 10),
            ("man", 12),
            ("proxy", 13),
            ("kmem", 15),
            ("dialout", 20),
            ("fax", 21),
            ("voice", 22),
            ("cdrom", 24),
            ("floppy", 25),
            ("tape", 26),
            ("sudo", 27),
            ("audio", 29),
            ("dip", 30),
            ("www-data", 33),
            ("backup", 34),
            ("operator", 37),
            ("list", 38),
            ("irc", 39),
            ("src", 40),
            ("gnats", 41),
            ("shadow", 42),
            ("utmp", 43),
            ("video", 44),
            ("sasl", 45),
            ("plugdev", 46),
            ("staff", 50),
            ("games", 60),
            ("users", 100),
            ("nogroup", 65534)
        ]
        
        # Create group file
        group_file = os.path.join(context.root, "etc", "group")
        with open(group_file, 'w') as f:
            for group_name, gid in default_groups:
                f.write(f"{group_name}:x:{gid}:\n")
        
        os.chmod(group_file, 0o644)
        logger.info(f"Created group file: {group_file}")
        
        # Create passwd file
        passwd_file = os.path.join(context.root, "etc", "passwd")
        with open(passwd_file, 'w') as f:
            # Root user
            f.write("root:x:0:0:root:/root:/bin/bash\n")
            # Debian user (for desktop systems)
            f.write("debian:x:1000:1000:Debian User:/home/debian:/bin/bash\n")
        
        os.chmod(passwd_file, 0o644)
        logger.info(f"Created passwd file: {passwd_file}")
        
        # Create shadow file
        shadow_file = os.path.join(context.root, "etc", "shadow")
        with open(shadow_file, 'w') as f:
            # Root user (locked)
            f.write("root:!:19131:0:99999:7:::\n")
            # Debian user (locked)
            f.write("debian:!:19131:0:99999:7:::\n")
        
        os.chmod(shadow_file, 0o640)
        logger.info(f"Created shadow file: {shadow_file}")
        
        # Create home directories
        home_dirs = ["root", "debian"]
        for user in home_dirs:
            user_home = os.path.join(context.root, "var", "home", user)
            os.makedirs(user_home, exist_ok=True)
            os.chmod(user_home, 0o700)
            logger.debug(f"Created home directory: {user_home}")
    
    def _setup_permissions(self, context) -> None:
        """
        Set up filesystem permissions.
        
        Args:
            context: osbuild context
        """
        logger.info("Setting up filesystem permissions")
        
        # Set permissions for essential directories
        permissions_map = {
            "/bin": 0o755,
            "/boot": 0o755,
            "/dev": 0o755,
            "/etc": 0o755,
            "/home": 0o755,
            "/lib": 0o755,
            "/lib64": 0o755,
            "/media": 0o755,
            "/mnt": 0o755,
            "/opt": 0o755,
            "/proc": 0o555,
            "/root": 0o700,
            "/run": 0o755,
            "/sbin": 0o755,
            "/srv": 0o755,
            "/sys": 0o555,
            "/tmp": 0o1777,
            "/usr": 0o755,
            "/var": 0o755,
            "/var/tmp": 0o1777,
            "/var/log": 0o755,
            "/var/cache": 0o755,
            "/var/lib": 0o755,
            "/var/spool": 0o755,
            "/var/lock": 0o755,
            "/var/run": 0o755
        }
        
        for path, mode in permissions_map.items():
            full_path = os.path.join(context.root, path.lstrip('/'))
            if os.path.exists(full_path):
                os.chmod(full_path, mode)
                logger.debug(f"Set permissions {oct(mode)} on {full_path}")
        
        # Set ownership for critical files
        ownership_map = {
            "/etc/passwd": (0, 0),  # root:root
            "/etc/group": (0, 0),   # root:root
            "/etc/shadow": (0, 0),  # root:root
            "/root": (0, 0),        # root:root
        }
        
        for path, (uid, gid) in ownership_map.items():
            full_path = os.path.join(context.root, path.lstrip('/'))
            if os.path.exists(full_path):
                try:
                    os.chown(full_path, uid, gid)
                    logger.debug(f"Set ownership {uid}:{gid} on {full_path}")
                except OSError as e:
                    # In test environments, we might not have permission to change ownership
                    # This is acceptable for testing purposes
                    logger.warning(f"Could not set ownership on {full_path}: {e}")
        
        logger.info("Filesystem permissions configured")
    
    def _configure_system_directories(self, context) -> None:
        """
        Configure system-specific directories.
        
        Args:
            context: osbuild context
        """
        logger.info("Configuring system directories")
        
        # Create systemd directories
        systemd_dirs = [
            "etc/systemd/system",
            "etc/systemd/user",
            "usr/lib/systemd/system",
            "usr/lib/systemd/user",
            "var/lib/systemd",
            "run/systemd"
        ]
        
        for directory in systemd_dirs:
            dir_path = os.path.join(context.root, directory)
            os.makedirs(dir_path, exist_ok=True)
            logger.debug(f"Created systemd directory: {dir_path}")
        
        # Create udev directories
        udev_dirs = [
            "etc/udev/rules.d",
            "usr/lib/udev/rules.d",
            "run/udev"
        ]
        
        for directory in udev_dirs:
            dir_path = os.path.join(context.root, directory)
            os.makedirs(dir_path, exist_ok=True)
            logger.debug(f"Created udev directory: {dir_path}")
        
        # Create APT directories
        apt_dirs = [
            "etc/apt/apt.conf.d",
            "etc/apt/sources.list.d",
            "var/lib/apt",
            "var/cache/apt"
        ]
        
        for directory in apt_dirs:
            dir_path = os.path.join(context.root, directory)
            os.makedirs(dir_path, exist_ok=True)
            logger.debug(f"Created APT directory: {dir_path}")
        
        # Create network configuration directories
        network_dirs = [
            "etc/network/interfaces.d",
            "etc/NetworkManager",
            "etc/NetworkManager/system-connections",
            "var/lib/NetworkManager"
        ]
        
        for directory in network_dirs:
            dir_path = os.path.join(context.root, directory)
            os.makedirs(dir_path, exist_ok=True)
            logger.debug(f"Created network directory: {dir_path}")
        
        # Create SSH directories
        ssh_dirs = [
            "etc/ssh",
            "var/lib/ssh",
            "var/run/sshd"
        ]
        
        for directory in ssh_dirs:
            dir_path = os.path.join(context.root, directory)
            os.makedirs(dir_path, exist_ok=True)
            logger.debug(f"Created SSH directory: {dir_path}")
        
        logger.info("System directories configured")


def main():
    """
    Main entry point for the Debian Filesystem stage.
    
    This function is called by osbuild when executing the stage.
    """
    import sys
    
    # Read options from stdin (osbuild passes options as JSON)
    options = json.load(sys.stdin)
    
    # Create and run the stage
    stage = DebianFilesystemStage(options)
    
    # Note: In a real osbuild stage, the context would be provided by osbuild
    # For now, this is a placeholder for testing
    class MockContext:
        def __init__(self, root):
            self.root = root
        
        def run(self, cmd):
            # Mock implementation for testing
            logger.info(f"Would run: {' '.join(cmd)}")
            return type('Result', (), {'returncode': 0, 'stdout': '', 'stderr': ''})()
    
    # For testing purposes
    if len(sys.argv) > 1 and sys.argv[1] == '--test':
        context = MockContext('/tmp/test-chroot')
        stage.run(context)
    else:
        # In real osbuild environment, context would be provided
        raise NotImplementedError("This stage must be run within osbuild")


if __name__ == "__main__":
    main()
