main.go

package main

import (
	"context"
	"fmt"
	"log"
	"os"
	"os/signal"
	"slices"
	"strings"
	"sync"
	"syscall"

	"github.com/fatih/color"
	"github.com/moby/moby/api/types/events"
	"github.com/moby/moby/client"
)

type LogState struct {
	LastTime string
	LastLine string
}

type Registry struct {
	mu           sync.Mutex
	active       map[string]context.CancelFunc
	lastStates   map[string]LogState
	Client       *Client
	Targets      []string
	targetColors map[string]*color.Color
}

func NewRegistry(client *Client, targets []string) *Registry {
	targetColors := make(map[string]*color.Color)
	{
		allColors := []*color.Color{
			color.New(color.FgCyan),
			color.New(color.FgBlue),
			color.New(color.FgMagenta),
			color.New(color.FgGreen),
			color.New(color.FgRed),
			color.New(color.FgYellow),
			color.New(color.FgCyan).Add(color.ReverseVideo),
			color.New(color.FgBlue).Add(color.ReverseVideo),
			color.New(color.FgMagenta).Add(color.ReverseVideo),
			color.New(color.FgGreen).Add(color.ReverseVideo),
			color.New(color.FgRed).Add(color.ReverseVideo),
			color.New(color.FgYellow).Add(color.ReverseVideo),
			color.New(color.FgCyan).Add(color.Bold).Add(color.Underline),
			color.New(color.FgBlue).Add(color.Bold).Add(color.Underline),
			color.New(color.FgMagenta).Add(color.Bold).Add(color.Underline),
			color.New(color.FgGreen).Add(color.Bold).Add(color.Underline),
			color.New(color.FgRed).Add(color.Bold).Add(color.Underline),
			color.New(color.FgYellow).Add(color.Bold).Add(color.Underline),
		}
		for i, v := range targets {
			targetColors[v] = allColors[i%len(allColors)]
		}
	}
	return &Registry{
		active:       make(map[string]context.CancelFunc),
		lastStates:   make(map[string]LogState),
		Client:       client,
		Targets:      targets,
		targetColors: targetColors,
	}
}

func main() {
	targets := os.Args[1:]
	if len(targets) == 0 {
		log.Fatal("Please specify at least one container name to tail")
	}

	slices.Sort(targets)
	targets = slices.Compact(targets)

	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stop()

	dockerClient := NewClient()

	reg := NewRegistry(dockerClient, targets)

	running, err := dockerClient.ListContainers(ctx)
	if err != nil {
		log.Println("error listing docker containers", err)
		os.Exit(1)
	}

	for _, c := range running {
		if reg.isTarget(c.Name, c.ComposeService) {
			reg.startTailing(ctx, c.ID, c.Name)
		}
	}

	eventChan, errChan := dockerClient.Events(ctx)

	for {
		select {
		case <-ctx.Done():
			return
		case err := <-errChan:
			log.Printf("Event Error: %v", err)
			return
		case msg := <-eventChan:
			reg.handleEvent(ctx, msg)
		}
	}
}

func (reg *Registry) handleEvent(ctx context.Context, msg events.Message) {
	name := strings.TrimPrefix(msg.Actor.Attributes["name"], "/")
	service := msg.Actor.Attributes["com.docker.compose.service"]

	if !reg.isTarget(name, service) {
		return
	}

	switch msg.Action {
	case "start":
		reg.startTailing(ctx, msg.Actor.ID, name)
	case "die", "kill", "stop":
		reg.stopTailing(name, name)
	}
}

func (reg *Registry) startTailing(ctx context.Context, id, name string) {
	reg.mu.Lock()
	defer reg.mu.Unlock()

	if _, exists := reg.active[name]; exists {
		return
	}

	displayName := name
	container, err := reg.Client.GetContainer(ctx, id)
	if err == nil && container.ComposeProject != "" && container.ComposeService != "" {
		displayName = shortenComposeName(container.ComposeProject, container.ComposeService)
	}

	tCtx, tCancel := context.WithCancel(ctx)
	reg.active[name] = tCancel

	state, hasState := reg.lastStates[name]

	go func() {
		defer reg.stopTailing(name, displayName)

		opts := client.ContainerLogsOptions{
			ShowStdout: true,
			ShowStderr: true,
			Follow:     true,
			Timestamps: true,
		}

		if hasState {
			opts.Since = state.LastTime
		} else {
			opts.Tail = "1000"
		}

		msgs, errs := reg.Client.ContainerLogs(tCtx, id, opts)

		for {
			select {
			case <-tCtx.Done():
				return
			case err := <-errs:
				if err != nil && err != context.Canceled {
					fmt.Printf("[%s] Stream closed: %v\n", displayName, err)
				}
				return
			case msg := <-msgs:
				parts := strings.SplitN(msg, " ", 2)
				if len(parts) < 2 {
					fmt.Printf("[%s] %s", displayName, msg)
					continue
				}

				timestamp := parts[0]
				content := parts[1]

				if hasState && msg == state.LastLine {
					continue
				}

				reg.mu.Lock()
				state = LogState{
					LastTime: timestamp,
					LastLine: msg,
				}
				reg.lastStates[name] = state
				hasState = true
				reg.mu.Unlock()

				var sb strings.Builder
				sb.WriteString(reg.targetColors[name].Sprint(displayName))
				sb.WriteString("┃")
				sb.WriteString(content)
				fmt.Printf(sb.String())
			}
		}
	}()
}

func (reg *Registry) stopTailing(name, displayName string) {
	reg.mu.Lock()
	defer reg.mu.Unlock()
	if cancel, exists := reg.active[name]; exists {
		cancel()
		delete(reg.active, name)
		fmt.Printf("--- Detached from %s ---\n", displayName)
	}
}

func (reg *Registry) isTarget(name, service string) bool {
	for _, t := range reg.Targets {
		if name == t || service == t {
			return true
		}
	}
	return false
}

func shortenComposeName(project, service string) string {
	parts := strings.Split(project, "-")
	var projectAbbr string
	for _, p := range parts {
		if len(p) > 0 {
			projectAbbr += string(p[0])
		}
	}
	if len(projectAbbr) < 2 && len(project) > 1 {
		lastPart := parts[len(parts)-1]
		needed := 2 - len(projectAbbr)
		if len(lastPart) > needed {
			projectAbbr += lastPart[len(lastPart)-needed:]
		} else {
			return fmt.Sprintf("%s-%s", project, service)
		}
	}

	return fmt.Sprintf("%s-%s", projectAbbr, service)
}