package main

import (
    "context"
    "flag"
    "fmt"
    "log"
    "os"
    "strings"

    "github.com/docker/docker/api/types"
    "github.com/docker/docker/client"
    "gopkg.in/yaml.v2"
)

func main() {
    // Command-line flags
    stackName := flag.String("s", "", "Specify the stack or project name to export.")
    outputFile := flag.String("o", "docker-compose.yml", "Specify the output file for the YAML.")
    flag.Parse()

    // Initialize Docker client
    cli, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation())
    if err != nil {
        log.Fatalf("Error initializing Docker client: %v", err)
    }

    ctx := context.Background()

    // Fetch list of all containers
    containers, err := cli.ContainerList(ctx, types.ContainerListOptions{})
    if err != nil {
        log.Fatalf("Error fetching container list: %v", err)
    }

    // Map to hold docker-compose configuration without version key
    composeConfig := make(map[string]interface{})
    composeConfig["services"] = make(map[string]interface{})

    // To store available stacks/projects if no specific stack is provided
    stacks := make(map[string]bool)

    // Process each container
    for _, container := range containers {
        project := container.Labels["com.docker.compose.project"]
        stackNamespace := container.Labels["com.docker.stack.namespace"]

        if *stackName == "" {
            // Gather available stack names
            if project != "" {
                stacks[project] = true
            } else if stackNamespace != "" {
                stacks[stackNamespace] = true
            }
        } else {
            // Export specific stack or project
            if project == *stackName || stackNamespace == *stackName {
                serviceName := strings.TrimPrefix(container.Names[0], "/")
                image := container.Image

                // Initialize service configuration
                serviceConfig := map[string]interface{}{
                    "image": image,
                }

                // Handle environment variables
                envVars := container.Config.Env
                if len(envVars) > 0 {
                    serviceConfig["environment"] = envVars
                }

                // Handle network settings
                networks := []string{}
                for network := range container.NetworkSettings.Networks {
                    networks = append(networks, network)
                }
                if len(networks) > 0 {
                    serviceConfig["networks"] = networks
                }

                // Add service configuration to compose config
                composeConfig["services"].(map[string]interface{})[serviceName] = serviceConfig
            }
        }
    }

    if *stackName == "" {
        // Print available stacks/projects
        fmt.Println("Available stacks or Docker Compose projects to export:")
        for stack := range stacks {
            fmt.Println("- " + stack)
        }
        return
    }

    // If specific stack/project is requested but not found
    if len(composeConfig["services"].(map[string]interface{})) == 0 {
        log.Fatalf("No services found for stack or project: %s", *stackName)
    }

    // Convert to YAML
    yamlData, err := yaml.Marshal(&composeConfig)
    if err != nil {
        log.Fatalf("Error generating YAML: %v", err)
    }

    // Write YAML to specified output file
    err = os.WriteFile(*outputFile, yamlData, 0644)
    if err != nil {
        log.Fatalf("Error writing to file: %v", err)
    }

    fmt.Printf("docker-compose.yml generated successfully in %s\n", *outputFile)
}