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", `{{json .Config.Env}} {{json .NetworkSettings.Networks}} {{json .Mounts}} {{json .HostConfig.PortBindings}} {{.HostConfig.RestartPolicy.Name}}`, containerID)
    var out bytes.Buffer
    cmd.Stdout = &out
    err := cmd.Run()
    if err != nil {
        return nil, nil, nil, nil, "", err
    }

    output := out.String()
    parts := strings.SplitN(output, " ", 5)
    envVars, networks, volumes, ports := []string{}, []string{}, []string{}, []string{}
    restartPolicy := ""

    if len(parts) > 0 && parts[0] != "null" {
        envVars = parseJSONStringArray(parts[0])
    }
    if len(parts) > 1 && parts[1] != "null" {
        networks = parseJSONNetworkNames(parts[1])
    }
    if len(parts) > 2 && parts[2] != "null" {
        volumes = parseJSONVolumeMounts(parts[2])
    }
    if len(parts) > 3 && parts[3] != "null" {
        ports = parseJSONPorts(parts[3])
    }
    if len(parts) > 4 {
        restartPolicy = strings.Trim(parts[4], `"`)
    }

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

func parseJSONStringArray(jsonStr string) []string {
    jsonStr = strings.Trim(jsonStr, "[]\"")
    if jsonStr == "" {
        return []string{}
    }
    return strings.Split(jsonStr, "\",\"")
}

func parseJSONNetworkNames(jsonStr string) []string {
    var networkNames []string
    networkData := strings.Split(jsonStr, ":")
    for _, data := range networkData {
        if strings.Contains(data, "\"Name\"") {
            parts := strings.Split(data, "\"")
            networkNames = append(networkNames, parts[1])
        }
    }
    return networkNames
}

func parseJSONVolumeMounts(jsonStr string) []string {
    var volumeMounts []string
    volumeData := strings.Split(jsonStr, "},{")
    for _, data := range volumeData {
        if strings.Contains(data, "\"Destination\"") {
            parts := strings.Split(data, "\"")
            volumeMounts = append(volumeMounts, parts[3]+":"+parts[7])
        }
    }
    return volumeMounts
}

func parseJSONPorts(jsonStr string) []string {
    var ports []string
    portData := strings.Split(jsonStr, "},{")
    for _, data := range portData {
        if strings.Contains(data, "\"HostPort\"") {
            hostPort := strings.Split(data, "\"")[3]
            containerPort := strings.Split(data, "\"")[7]
            ports = append(ports, hostPort+":"+containerPort)
        }
    }
    return ports
}

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
}

