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"]

        // Use project name if available; otherwise, use 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)
    }

    // Determine if this is a Docker Swarm stack or a Docker Compose project
    isSwarm := isDockerSwarmStack(*stackName, containers)

    if isSwarm {
        // Handle Docker Swarm stack export
        err = exportSwarmStackToComposeFile(*stackName, *outputFile)
        if err != nil {
            log.Fatalf("Error exporting Docker Swarm stack %s: %v", *stackName, err)
        }
        fmt.Printf("docker-compose.yml generated successfully for Docker Swarm stack '%s' in %s\n", *stackName, *outputFile)
    } else {
        // Handle Docker Compose project export
        err = exportComposeProjectToComposeFile(*stackName, containers, *outputFile)
        if err != nil {
            log.Fatalf("Error exporting Docker Compose project %s: %v", *stackName, err)
        }
        fmt.Printf("docker-compose.yml generated successfully for Docker Compose project '%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())

        // Check that there are at least 3 fields (ID, Name, Image)
        if len(fields) < 3 {
            log.Printf("Warning: Skipping incomplete line: %s", scanner.Text())
            continue
        }

        id := fields[0]
        name := fields[1]
        image := fields[2]

        // Handle cases where project or stackNamespace might be missing
        project := ""
        stackNamespace := ""
        if len(fields) > 3 {
            project = fields[3]
        }
        if len(fields) > 4 {
            stackNamespace = fields[4]
        }

        labels := map[string]string{
            "com.docker.compose.project": project,
            "com.docker.stack.namespace":  stackNamespace,
        }

        // Run 'docker inspect' to get detailed container information.
        env, networks, err := inspectContainer(id)
        if err != nil {
            log.Printf("Warning: Error inspecting container %s: %v", id, err)
            continue
        }

        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
    }

    // Split output into environment variables and network names
    parts := strings.Split(out.String(), "  ")
    if len(parts) < 2 {
        // Handle cases where no networks are listed
        parts = append(parts, "")
    }
    envVars := strings.Fields(parts[0])
    networks := strings.Fields(parts[1])

    return envVars, networks, nil
}

// isDockerSwarmStack checks if the stack is a Docker Swarm stack based on the container names.
func isDockerSwarmStack(stackName string, containers []Container) bool {
    for _, container := range containers {
        if strings.Contains(container.Name, ".") {
            return true
        }
    }
    return false
}

// exportSwarmStackToComposeFile exports a Docker Swarm stack to a docker-compose.yml file.
func exportSwarmStackToComposeFile(stackName string, outputFile string) error {
    // Fetch stack details using `docker stack services` and construct the docker-compose.yml
    cmd := exec.Command("docker", "stack", "services", stackName, "--format", "{{.Name}} {{.Image}} {{.Ports}} {{.Replicas}}")
    var out bytes.Buffer
    cmd.Stdout = &out
    err := cmd.Run()
    if err != nil {
        return fmt.Errorf("error fetching Docker Swarm stack services: %v", err)
    }

    composeConfig := make(map[string]interface{})
    composeConfig["services"] = make(map[string]interface{})

    scanner := bufio.NewScanner(&out)
    for scanner.Scan() {
        fields := strings.Fields(scanner.Text())
        if len(fields) < 2 {
            continue
        }

        serviceName := fields[0]
        image := fields[1]

        serviceConfig := map[string]interface{}{
            "image": image,
        }

        // Handle additional fields like ports and replicas
        if len(fields) > 2 {
            // Parse ports and replicas if available
            ports := fields[2]
            replicas := fields[3]
            serviceConfig["deploy"] = map[string]interface{}{
                "replicas": replicas,
            }
            if ports != "" {
                serviceConfig["ports"] = []string{ports}
            }
        }

        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 Docker Swarm stack %s: %v", stackName, 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
}

// exportComposeProjectToComposeFile exports a Docker Compose project to a docker-compose.yml file.
func exportComposeProjectToComposeFile(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 Docker Compose project %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
}