package main

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

    "gopkg.in/yaml.v2"
)

func main() {
    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()

    containers, err := listContainers()
    if err != nil {
        log.Fatalf("Error fetching container list: %v", err)
    }

    stacks := make(map[string][]Container)

    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 == "" {
        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
    }

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

    isSwarm := isDockerSwarmStack(*stackName, containers)

    if isSwarm {
        fmt.Printf("Warning: The stack '%s' appears to be a Docker Swarm stack. Export functionality for Swarm stacks may not provide complete configurations.\n", *stackName)
        fmt.Println("Skipping export for Docker Swarm stacks. Please use `docker stack` commands directly for managing Swarm stacks.")
        return
    } else {
        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)
    }
}

type Container struct {
    ID            string
    Name          string
    Image         string
    Labels        map[string]string
    Env           []string
    Networks      []string
    Volumes       []string
    Ports         []string
    RestartPolicy string
}

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 {
            log.Printf("Warning: Skipping incomplete line: %s", scanner.Text())
            continue
        }

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

        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,
        }

        env, networks, volumes, ports, restartPolicy, 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,
            Volumes:       volumes,
            Ports:         ports,
            RestartPolicy: restartPolicy,
        })
    }

    return containers, scanner.Err()
}

func inspectContainer(containerID string) ([]string, []string, []string, []string, string, error) {
    cmd := exec.Command("docker", "inspect", "--format", `{{range .Config.Env}}{{println .}}{{end}}{{range .NetworkSettings.Networks}}{{println .}}{{end}}{{range .Mounts}}{{println .}}{{end}}{{range $p, $conf := .HostConfig.PortBindings}}{{println $p}}{{end}}{{.HostConfig.RestartPolicy.Name}}`, containerID)
    var out bytes.Buffer
    cmd.Stdout = &out
    err := cmd.Run()
    if err != nil {
        return nil, nil, nil, nil, "", err
    }

    lines := strings.Split(out.String(), "\n")

    var envVars, networks, volumes, ports []string
    restartPolicy := ""

    for _, line := range lines {
        if strings.Contains(line, "=") {
            envVars = append(envVars, line)
        } else if strings.Contains(line, "Name:") && strings.Contains(line, "Driver:") {
            networks = append(networks, strings.TrimSpace(line))
        } else if strings.Contains(line, "Source:") && strings.Contains(line, "Destination:") {
            volume := strings.Split(line, "Destination:")[1]
            volumes = append(volumes, volume)
        } else if strings.Contains(line, "/tcp") || strings.Contains(line, "/udp") {
            ports = append(ports, line)
        } else if line != "" {
            restartPolicy = line
        }
    }

    return envVars, networks, volumes, ports, restartPolicy, nil
}

func isDockerSwarmStack(stackName string, containers []Container) bool {
    for _, container := range containers {
        if strings.Contains(container.Name, ".") {
            return true
        }
    }
    return false
}

func exportComposeProjectToComposeFile(stack string, containers []Container, outputFile string) error {
    composeConfig := make(map[string]interface{})
    composeConfig["services"] = make(map[string]interface{})

    for _, container := range containers {
        serviceName := container.Name
        image := container.Image

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

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

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

        if len(container.Volumes) > 0 {
            serviceConfig["volumes"] = container.Volumes
        }

        if len(container.Ports) > 0 {
            serviceConfig["ports"] = container.Ports
        }

        if container.RestartPolicy != "" {
            serviceConfig["restart"] = container.RestartPolicy
        }

        composeConfig["services"].(map[string]interface{})[serviceName] = serviceConfig
    }

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

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

    return nil
}