#!/usr/bin/env python3

import pytest
import tempfile
import os

def test_sources_stage_core_logic():
    """Test the core logic of the sources stage"""
    
    def main(tree, options):
        """Configure APT sources.list for the target filesystem"""
        
        # Get options
        sources = options.get("sources", [])
        suite = options.get("suite", "trixie")
        mirror = options.get("mirror", "https://deb.debian.org/debian")
        components = options.get("components", ["main"])
        
        # Default sources if none provided
        if not sources:
            sources = [
                {
                    "type": "deb",
                    "uri": mirror,
                    "suite": suite,
                    "components": components
                }
            ]
        
        # Create sources.list.d directory
        sources_dir = os.path.join(tree, "etc", "apt", "sources.list.d")
        os.makedirs(sources_dir, exist_ok=True)
        
        # Clear existing sources.list
        sources_list = os.path.join(tree, "etc", "apt", "sources.list")
        if os.path.exists(sources_list):
            os.remove(sources_list)
        
        # Create new sources.list
        with open(sources_list, "w") as f:
            for source in sources:
                source_type = source.get("type", "deb")
                uri = source.get("uri", mirror)
                source_suite = source.get("suite", suite)
                source_components = source.get("components", components)
                
                # Handle different source types
                if source_type == "deb":
                    f.write(f"{source_type} {uri} {source_suite} {' '.join(source_components)}\n")
                elif source_type == "deb-src":
                    f.write(f"{source_type} {uri} {source_suite} {' '.join(source_components)}\n")
                elif source_type == "deb-ports":
                    f.write(f"{source_type} {uri} {source_suite} {' '.join(source_components)}\n")
        
        print(f"APT sources configured for {suite}")
        return 0
    
    # Test with custom options
    with tempfile.TemporaryDirectory() as temp_dir:
        os.makedirs(os.path.join(temp_dir, "etc", "apt"), exist_ok=True)
        
        result = main(temp_dir, {
            "suite": "trixie",
            "mirror": "https://deb.debian.org/debian",
            "components": ["main", "contrib"]
        })
        
        assert result == 0
        
        # Check that sources.list was created
        sources_list = os.path.join(temp_dir, "etc", "apt", "sources.list")
        assert os.path.exists(sources_list)
        
        # Check content
        with open(sources_list, 'r') as f:
            content = f.read()
            assert "deb https://deb.debian.org/debian trixie main contrib" in content

def test_sources_stage_defaults():
    """Test the sources stage with default options"""
    
    def main(tree, options):
        """Configure APT sources.list for the target filesystem"""
        
        # Get options
        sources = options.get("sources", [])
        suite = options.get("suite", "trixie")
        mirror = options.get("mirror", "https://deb.debian.org/debian")
        components = options.get("components", ["main"])
        
        # Default sources if none provided
        if not sources:
            sources = [
                {
                    "type": "deb",
                    "uri": mirror,
                    "suite": suite,
                    "components": components
                }
            ]
        
        # Create sources.list.d directory
        sources_dir = os.path.join(tree, "etc", "apt", "sources.list.d")
        os.makedirs(sources_dir, exist_ok=True)
        
        # Clear existing sources.list
        sources_list = os.path.join(tree, "etc", "apt", "sources.list")
        if os.path.exists(sources_list):
            os.remove(sources_list)
        
        # Create new sources.list
        with open(sources_list, "w") as f:
            for source in sources:
                source_type = source.get("type", "deb")
                uri = source.get("uri", mirror)
                source_suite = source.get("suite", suite)
                source_components = source.get("components", components)
                
                # Handle different source types
                if source_type == "deb":
                    f.write(f"{source_type} {uri} {source_suite} {' '.join(source_components)}\n")
                elif source_type == "deb-src":
                    f.write(f"{source_type} {uri} {source_suite} {' '.join(source_components)}\n")
                elif source_type == "deb-ports":
                    f.write(f"{source_type} {uri} {source_suite} {' '.join(source_components)}\n")
        
        print(f"APT sources configured for {suite}")
        return 0
    
    with tempfile.TemporaryDirectory() as temp_dir:
        os.makedirs(os.path.join(temp_dir, "etc", "apt"), exist_ok=True)
        
        result = main(temp_dir, {})
        assert result == 0
        
        sources_list = os.path.join(temp_dir, "etc", "apt", "sources.list")
        assert os.path.exists(sources_list)
        
        with open(sources_list, 'r') as f:
            content = f.read()
            assert "deb https://deb.debian.org/debian trixie main" in content

if __name__ == "__main__":
    pytest.main([__file__])