#!/usr/bin/env python3
"""
Container build utilities for deb-bootc-image-builder.

This module provides utilities for building and testing container images,
including:
- Container image building
- Container validation
- Debian-specific container features
"""

import os
import subprocess
import tempfile
import shutil
import json
import logging
from typing import Dict, List, Any, Optional, Tuple
from pathlib import Path

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


class ContainerBuilder:
    """Build and manage Debian container images."""

    def __init__(self, work_dir: str):
        """Initialize container builder."""
        self.work_dir = work_dir
        self.container_dir = os.path.join(work_dir, "containers")
        os.makedirs(self.container_dir, exist_ok=True)

    def build_debian_container(self, 
                             base_image: str = "debian:bookworm",
                             packages: Optional[List[str]] = None,
                             customizations: Optional[Dict[str, Any]] = None) -> str:
        """Build a Debian container image."""
        if packages is None:
            packages = [
                "linux-image-amd64",
                "systemd",
                "ostree",
                "grub-efi-amd64",
                "initramfs-tools"
            ]
        
        if customizations is None:
            customizations = {}
        
        # Create Containerfile
        containerfile_path = self._create_containerfile(
            base_image, packages, customizations
        )
        
        # Build container
        image_name = f"debian-bootc-{os.path.basename(work_dir)}"
        image_tag = "latest"
        
        build_result = self._build_container(containerfile_path, image_name, image_tag)
        
        if build_result["success"]:
            logger.info(f"Container built successfully: {image_name}:{image_tag}")
            return f"{image_name}:{image_tag}"
        else:
            raise RuntimeError(f"Container build failed: {build_result['error']}")

    def _create_containerfile(self, 
                            base_image: str, 
                            packages: List[str], 
                            customizations: Dict[str, Any]) -> str:
        """Create a Containerfile for building."""
        containerfile_content = f"""FROM {base_image}

# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive

# Install essential packages
RUN apt-get update && apt-get install -y \\
    {' \\\n    '.join(packages)} \\
    && apt-get clean \\
    && rm -rf /var/lib/apt/lists/*

# Set up OSTree configuration
RUN mkdir -p /etc/ostree \\
    && echo '[core]' > /etc/ostree/ostree.conf \\
    && echo 'mode=bare-user-only' >> /etc/ostree/ostree.conf

# Configure system identification
RUN echo 'PRETTY_NAME="Debian Bootc Image"' > /etc/os-release \\
    && echo 'NAME="Debian"' >> /etc/os-release \\
    && echo 'VERSION="13"' >> /etc/os-release \\
    && echo 'ID=debian' >> /etc/os-release \\
    && echo 'ID_LIKE=debian' >> /etc/os-release \\
    && echo 'VERSION_ID="13"' >> /etc/os-release

# Set up /home -> /var/home symlink for immutable architecture
RUN ln -sf /var/home /home

# Apply customizations
"""
        
        # Add customizations
        if "users" in customizations:
            for user in customizations["users"]:
                containerfile_content += f"RUN useradd -m -G sudo {user['name']} \\\n"
                if "password" in user:
                    containerfile_content += f"    && echo '{user['name']}:{user['password']}' | chpasswd \\\n"
                containerfile_content += f"    && echo '{user['name']} ALL=(ALL) NOPASSWD:ALL' >> /etc/sudoers\n"
        
        if "services" in customizations:
            for service in customizations["services"]:
                containerfile_content += f"RUN systemctl enable {service}\n"
        
        # Add labels
        containerfile_content += """
# Add labels
LABEL org.opencontainers.image.title="Debian Bootc Image"
LABEL org.opencontainers.image.description="Debian-based bootc image"
LABEL org.opencontainers.image.vendor="Debian Project"
LABEL org.opencontainers.image.version="13"
LABEL com.debian.bootc="true"
LABEL ostree.bootable="true"
"""
        
        # Write Containerfile
        containerfile_path = os.path.join(self.container_dir, "Containerfile")
        with open(containerfile_path, 'w') as f:
            f.write(containerfile_content)
        
        return containerfile_path

    def _build_container(self, containerfile_path: str, image_name: str, image_tag: str) -> Dict[str, Any]:
        """Build container using podman."""
        try:
            cmd = [
                "podman", "build",
                "-f", containerfile_path,
                "-t", f"{image_name}:{image_tag}",
                "."
            ]
            
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                cwd=self.container_dir
            )
            
            if result.returncode == 0:
                return {
                    "success": True,
                    "image": f"{image_name}:{image_tag}",
                    "output": result.stdout
                }
            else:
                return {
                    "success": False,
                    "error": result.stderr,
                    "returncode": result.returncode
                }
        
        except Exception as e:
            return {
                "success": False,
                "error": str(e)
            }

    def validate_container(self, image_name: str, image_tag: str) -> Dict[str, Any]:
        """Validate a built container image."""
        try:
            # Check if image exists
            cmd = ["podman", "image", "exists", f"{image_name}:{image_tag}"]
            result = subprocess.run(cmd, capture_output=True, text=True)
            
            if result.returncode != 0:
                return {
                    "valid": False,
                    "error": f"Image {image_name}:{image_tag} does not exist"
                }
            
            # Inspect image
            cmd = ["podman", "inspect", f"{image_name}:{image_tag}"]
            result = subprocess.run(cmd, capture_output=True, text=True)
            
            if result.returncode != 0:
                return {
                    "valid": False,
                    "error": f"Failed to inspect image: {result.stderr}"
                }
            
            # Parse inspection output
            try:
                image_info = json.loads(result.stdout)
                if isinstance(image_info, list):
                    image_info = image_info[0]
                
                # Validate labels
                labels = image_info.get("Labels", {})
                required_labels = ["com.debian.bootc", "ostree.bootable"]
                
                validation_result = {
                    "valid": True,
                    "labels": labels,
                    "architecture": image_info.get("Architecture", "unknown"),
                    "os": image_info.get("Os", "unknown"),
                    "size": image_info.get("Size", 0)
                }
                
                for label in required_labels:
                    if label not in labels:
                        validation_result["valid"] = False
                        validation_result["error"] = f"Missing required label: {label}"
                        break
                
                return validation_result
            
            except json.JSONDecodeError as e:
                return {
                    "valid": False,
                    "error": f"Failed to parse image inspection: {e}"
                }
        
        except Exception as e:
            return {
                "valid": False,
                "error": f"Validation failed: {e}"
            }

    def run_container_test(self, image_name: str, image_tag: str) -> Dict[str, Any]:
        """Run basic tests on a container image."""
        try:
            # Test container startup
            cmd = [
                "podman", "run", "--rm",
                f"{image_name}:{image_tag}",
                "echo", "Container test successful"
            ]
            
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=30
            )
            
            if result.returncode == 0:
                return {
                    "success": True,
                    "test": "container_startup",
                    "output": result.stdout.strip()
                }
            else:
                return {
                    "success": False,
                    "test": "container_startup",
                    "error": result.stderr,
                    "returncode": result.returncode
                }
        
        except subprocess.TimeoutExpired:
            return {
                "success": False,
                "test": "container_startup",
                "error": "Container startup timed out"
            }
        except Exception as e:
            return {
                "success": False,
                "test": "container_startup",
                "error": str(e)
            }

    def cleanup_container(self, image_name: str, image_tag: str) -> bool:
        """Clean up a container image."""
        try:
            cmd = ["podman", "rmi", f"{image_name}:{image_tag}"]
            result = subprocess.run(cmd, capture_output=True, text=True)
            
            if result.returncode == 0:
                logger.info(f"Container {image_name}:{image_tag} cleaned up successfully")
                return True
            else:
                logger.warning(f"Failed to clean up container: {result.stderr}")
                return False
        
        except Exception as e:
            logger.error(f"Error during container cleanup: {e}")
            return False


def create_test_container(work_dir: str, 
                         base_image: str = "debian:bookworm",
                         packages: Optional[List[str]] = None) -> Tuple[str, ContainerBuilder]:
    """Create a test container and return the builder instance."""
    builder = ContainerBuilder(work_dir)
    
    if packages is None:
        packages = [
            "linux-image-amd64",
            "systemd",
            "ostree"
        ]
    
    image_name = builder.build_debian_container(base_image, packages)
    return image_name, builder


def test_container_build_workflow(work_dir: str):
    """Test the complete container build workflow."""
    # Create container builder
    builder = ContainerBuilder(work_dir)
    
    # Define test packages
    test_packages = [
        "linux-image-amd64",
        "systemd",
        "ostree",
        "grub-efi-amd64",
        "initramfs-tools"
    ]
    
    # Define customizations
    customizations = {
        "users": [
            {"name": "testuser", "password": "testpass"}
        ],
        "services": ["systemd-networkd", "dbus"]
    }
    
    try:
        # Build container
        image_name = builder.build_debian_container(
            base_image="debian:bookworm",
            packages=test_packages,
            customizations=customizations
        )
        
        # Validate container
        validation_result = builder.validate_container(image_name, "latest")
        assert validation_result["valid"], f"Container validation failed: {validation_result.get('error', 'Unknown error')}"
        
        # Run container test
        test_result = builder.run_container_test(image_name, "latest")
        assert test_result["success"], f"Container test failed: {test_result.get('error', 'Unknown error')}"
        
        logger.info("Container build workflow test completed successfully")
        
        # Cleanup
        builder.cleanup_container(image_name, "latest")
        
        return True
    
    except Exception as e:
        logger.error(f"Container build workflow test failed: {e}")
        return False


if __name__ == "__main__":
    # Test the container builder
    work_dir = tempfile.mkdtemp()
    try:
        success = test_container_build_workflow(work_dir)
        if success:
            print("Container build workflow test passed!")
        else:
            print("Container build workflow test failed!")
    finally:
        shutil.rmtree(work_dir)
