import os
import yaml
import csv
from dotenv import load_dotenv

def load_env_vars(env_file):
    """Load environment variables from a .env file."""
    if os.path.exists(env_file):
        load_dotenv(env_file)

def expand_env_vars(value):
    """Expand environment variables in a string."""
    return os.path.expandvars(value)

def parse_docker_compose(file_path):
    """Parse a docker-compose.yml file and extract volume information."""
    with open(file_path, 'r') as file:
        compose_data = yaml.safe_load(file)
    
    volumes = []
    if compose_data and 'services' in compose_data:
        for service_name, service_config in compose_data['services'].items():
            if 'volumes' in service_config:
                for volume in service_config['volumes']:
                    if isinstance(volume, str):
                        # Handle string format: "host_path:container_path"
                        parts = volume.split(':')
                        if len(parts) == 2:
                            host_path, container_path = parts
                            volumes.append((host_path, container_path))
                        elif len(parts) == 3:
                            # Handle cases like "host_path:container_path:ro"
                            host_path, container_path, _ = parts
                            volumes.append((host_path, container_path))
                    elif isinstance(volume, dict):
                        # Handle dictionary format: { 'source': 'host_path', 'target': 'container_path' }
                        if 'source' in volume and 'target' in volume:
                            volumes.append((volume['source'], volume['target']))
    return volumes

def get_volume_type(host_path):
    """Determine if the volume is a bind mount or a named volume."""
    if host_path.startswith(('.', '/', '~')) or '${' in host_path:
        return "bind"
    else:
        return "named"

def traverse_directory(root_dir):
    """Traverse the directory and find all docker-compose.yml files."""
    results = []
    for dirpath, _, filenames in os.walk(root_dir):
        if 'docker-compose.yml' in filenames:
            compose_file = os.path.join(dirpath, 'docker-compose.yml')
            env_file = os.path.join(dirpath, '.env')
            load_env_vars(env_file)
            volumes = parse_docker_compose(compose_file)
            for host_path, container_path in volumes:
                expanded_host_path = expand_env_vars(host_path)
                volume_type = get_volume_type(expanded_host_path)
                results.append((dirpath, expanded_host_path, container_path, volume_type))
    return results

def write_to_csv(data, output_file):
    """Write the volume data to a CSV file."""
    headers = ["Directory", "Volume (Host Path)", "Bind Point (Container Path)", "Volume Type"]
    with open(output_file, mode='w', newline='') as file:
        writer = csv.writer(file)
        writer.writerow(headers)  # Write the header row
        writer.writerows(data)  # Write the data rows

def main():
    root_directory = input("Enter the root directory to traverse: ")
    if not os.path.isdir(root_directory):
        print("Invalid directory path.")
        return
    
    output_csv = input("Enter the output CSV file path (e.g., output.csv): ")
    
    # Traverse the directory and get volume information
    volume_data = traverse_directory(root_directory)
    
    # Write the results to a CSV file
    if volume_data:
        write_to_csv(volume_data, output_csv)
        print(f"Volume data has been written to {output_csv}")
    else:
        print("No docker-compose.yml files with volumes found.")

if __name__ == "__main__":
    main()