package main

import (
    "bufio"
    "bytes"
    "flag"
    "fmt"
    "log"
    "os"
    "os/exec"
    "strings"

    "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()

    // Fetch list of all containers with their stack or project labels
    containers, err := listContainers()
    if err != nil {
        log.Fatalf("Error fetching container list: %v", err)
    }

    // Map to group containers by stack or project
    stacks := make(map[string][]Container)

    // Group containers by Docker Compose project or Docker Swarm stack
    for _, container := range containers {
        project := container.Labels["com.docker.compose.project"]
        stackNamespace := container.Labels["com.docker.stack.namespace"]
        stackKey := project
        if stackNamespace != "" {
            stackKey = stackNamespace
        }

        if stackKey != "" {
            stacks[stackKey] = append(stacks[stackKey], container)
        }
    }

    if *stackName == "" {
        // No stack specified, list available stacks
        if len(stacks) == 0 {
            fmt.Println("No Docker Compose projects or Docker Swarm stacks found.")
        } else {
            fmt.Println("Available stacks or Docker Compose projects to export:")
            for stack := range stacks {
                fmt.Println("- " + stack)
            }
        }
        return
    }

    // Stack specified, check if it exists
    containers, exists := stacks[*stackName]
    if !exists {
        log.Fatalf("Stack or project '%s' not found.", *stackName)
    }

    // Export the specified stack/project to a docker-compose.yml file
    err = exportStackToComposeFile(*stackName, containers, *outputFile)
    if err != nil {
        log.Fatalf("Error exporting stack %s: %v", *stackName, err)
    }
    fmt.Printf("docker-compose.yml generated successfully for stack '%s' in %s\n", *stackName, *outputFile)
}

// Container represents basic Docker container information.
type Container struct {
    ID       string
    Name     string
    Image    string
    Labels   map[string]string
    Env      []string
    Networks []string
}

// listContainers runs 'docker ps' and 'docker inspect' to get a list of all containers.
func listContainers() ([]Container, error) {
    cmd := exec.Command("docker", "ps", "--format", "{{.ID}} {{.Names}} {{.Image}} {{.Label \"com.docker.compose.project\"}} {{.Label \"com.docker.stack.namespace\"}}")
    var out bytes.Buffer
    cmd.Stdout = &out
    err := cmd.Run()
    if err != nil {
        return nil, err
    }

    scanner := bufio.NewScanner(&out)
    var containers []Container

    for scanner.Scan() {
        fields := strings.Fields(scanner.Text())
        if len(fields) < 3 {
            continue
        }

        id := fields[0]
        name := fields[1]
        image := fields[2]
        labels := map[string]string{
            "com.docker.compose.project": fields[3],
            "com.docker.stack.namespace":  fields[4],
        }

        // Run 'docker inspect' to get detailed container information.
        env, networks, err := inspectContainer(id)
        if err != nil {
            return nil, err
        }

        containers = append(containers, Container{
            ID:       id,
            Name:     name,
            Image:    image,
            Labels:   labels,
            Env:      env,
            Networks: networks,
        })
    }

    return containers, scanner.Err()
}

// inspectContainer runs 'docker inspect' and parses environment variables and networks.
func inspectContainer(containerID string) ([]string, []string, error) {
    cmd := exec.Command("docker", "inspect", "--format", "{{range .Config.Env}}{{printf \"%s \" .}}{{end}} {{range $key, $value := .NetworkSettings.Networks}}{{$key}} {{end}}", containerID)
    var out bytes.Buffer
    cmd.Stdout = &out
    err := cmd.Run()
    if err != nil {
        return nil, nil, err
    }

    parts := strings.Split(out.String(), "  ")
    envVars := strings.Fields(parts[0])
    networks := strings.Fields(parts[1])

    return envVars, networks, nil
}

// exportStackToComposeFile generates a docker-compose.yml for a given stack or project.
func exportStackToComposeFile(stack string, containers []Container, outputFile string) error {
    // Map to hold docker-compose configuration
    composeConfig := make(map[string]interface{})
    composeConfig["services"] = make(map[string]interface{})

    // Process each container in the stack
    for _, container := range containers {
        serviceName := container.Name
        image := container.Image

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

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

        // Handle network settings
        if len(container.Networks) > 0 {
            serviceConfig["networks"] = container.Networks
        }

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

    // Convert to YAML
    yamlData, err := yaml.Marshal(&composeConfig)
    if err != nil {
        return fmt.Errorf("error generating YAML for stack %s: %v", stack, err)
    }

    // Write YAML to a file
    err = os.WriteFile(outputFile, yamlData, 0644)
    if err != nil {
        return fmt.Errorf("error writing to file %s: %v", outputFile, err)
    }

    return nil
}