#!/usr/bin/env python3
"""
Generate apt-ostree treefiles from variant definitions
"""

import os
import yaml
import argparse
from pathlib import Path
from typing import Dict, List, Any

class VariantGenerator:
    """Generate apt-ostree treefiles from variant definitions"""
    
    def __init__(self, config_dir: str):
        self.config_dir = Path(config_dir)
        self.variants_dir = self.config_dir / "variants"
        self.package_groups_dir = self.config_dir / "package-groups"
        self.treefiles_dir = self.config_dir / "treefiles"
        
        # Ensure output directory exists
        self.treefiles_dir.mkdir(exist_ok=True)
    
    def load_variant(self, variant_file: str) -> Dict[str, Any]:
        """Load a variant definition file"""
        variant_path = self.variants_dir / variant_file
        with open(variant_path, 'r') as f:
            return yaml.safe_load(f)
    
    def load_package_group(self, group_file: str) -> Dict[str, Any]:
        """Load a package group definition file"""
        group_path = self.package_groups_dir / group_file
        with open(group_path, 'r') as f:
            return yaml.safe_load(f)
    
    def flatten_packages(self, variant: Dict[str, Any]) -> List[str]:
        """Flatten all package lists into a single list"""
        packages = []
        
        # Add base packages
        if 'base_packages' in variant:
            packages.extend(variant['base_packages'])
        
        # Add all other package categories
        for key, value in variant.items():
            if key.endswith('_packages') and key != 'base_packages':
                if isinstance(value, list):
                    packages.extend(value)
        
        # Remove duplicates while preserving order
        seen = set()
        unique_packages = []
        for pkg in packages:
            if pkg not in seen:
                seen.add(pkg)
                unique_packages.append(pkg)
        
        return unique_packages
    
    def generate_treefile(self, variant: Dict[str, Any]) -> Dict[str, Any]:
        """Generate an apt-ostree treefile from variant definition"""
        
        # Get all packages
        packages = self.flatten_packages(variant)
        
        # Get excluded packages
        excluded = variant.get('excluded_packages', [])
        
        # Remove excluded packages
        final_packages = [pkg for pkg in packages if pkg not in excluded]
        
        # Build treefile structure
        treefile = {
            'ref': variant['ostree']['ref'],
            'repos': variant['repositories'],
            'packages': {
                'include': final_packages
            }
        }
        
        # Add system configuration if present
        if 'system_config' in variant:
            treefile['system'] = variant['system_config']
        
        return treefile
    
    def save_treefile(self, treefile: Dict[str, Any], variant_name: str) -> str:
        """Save treefile to disk"""
        filename = f"{variant_name}.yaml"
        output_path = self.treefiles_dir / filename
        
        with open(output_path, 'w') as f:
            yaml.dump(treefile, f, default_flow_style=False, indent=2)
        
        return str(output_path)
    
    def generate_all_variants(self) -> List[str]:
        """Generate treefiles for all variants"""
        generated_files = []
        
        # Find all variant files
        variant_files = list(self.variants_dir.glob("*.yaml"))
        
        for variant_file in variant_files:
            print(f"Processing variant: {variant_file.name}")
            
            try:
                # Load variant
                variant = self.load_variant(variant_file.name)
                
                # Generate treefile
                treefile = self.generate_treefile(variant)
                
                # Save treefile
                variant_name = variant['variant']['name']
                output_path = self.save_treefile(treefile, variant_name)
                
                print(f"  Generated: {output_path}")
                generated_files.append(output_path)
                
            except Exception as e:
                print(f"  Error processing {variant_file.name}: {e}")
        
        return generated_files
    
    def validate_treefile(self, treefile_path: str) -> bool:
        """Validate a generated treefile"""
        try:
            with open(treefile_path, 'r') as f:
                treefile = yaml.safe_load(f)
            
            # Check required fields
            required_fields = ['ref', 'repos', 'packages']
            for field in required_fields:
                if field not in treefile:
                    print(f"  ❌ Missing required field: {field}")
                    return False
            
            # Check packages structure
            if 'include' not in treefile['packages']:
                print(f"  ❌ Missing packages.include field")
                return False
            
            # Check repositories structure
            for repo in treefile['repos']:
                if 'name' not in repo or 'url' not in repo:
                    print(f"  ❌ Invalid repository structure: {repo}")
                    return False
            
            print(f"  ✅ Valid treefile: {treefile_path}")
            return True
            
        except Exception as e:
            print(f"  ❌ Error validating {treefile_path}: {e}")
            return False

def main():
    parser = argparse.ArgumentParser(description="Generate apt-ostree treefiles from variants")
    parser.add_argument("--config-dir", default=".", help="Configuration directory")
    parser.add_argument("--validate", action="store_true", help="Validate generated treefiles")
    parser.add_argument("--variant", help="Generate treefile for specific variant only")
    
    args = parser.parse_args()
    
    generator = VariantGenerator(args.config_dir)
    
    if args.variant:
        # Generate single variant
        variant_file = f"{args.variant}.yaml"
        if not (generator.variants_dir / variant_file).exists():
            print(f"Variant file not found: {variant_file}")
            return 1
        
        variant = generator.load_variant(variant_file)
        treefile = generator.generate_treefile(variant)
        variant_name = variant['variant']['name']
        output_path = generator.save_treefile(treefile, variant_name)
        
        print(f"Generated: {output_path}")
        
        if args.validate:
            generator.validate_treefile(output_path)
    else:
        # Generate all variants
        print("Generating treefiles for all variants...")
        generated_files = generator.generate_all_variants()
        
        print(f"\nGenerated {len(generated_files)} treefiles:")
        for file_path in generated_files:
            print(f"  {file_path}")
        
        if args.validate:
            print("\nValidating treefiles...")
            valid_count = 0
            for file_path in generated_files:
                if generator.validate_treefile(file_path):
                    valid_count += 1
            
            print(f"\nValidation complete: {valid_count}/{len(generated_files)} treefiles valid")

if __name__ == "__main__":
    main()