package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"os/exec"
	"path/filepath"
	"regexp"
	"sort"
	"strconv"
	"strings"
)

// Facts about a game server image, read without starting the game: Docker metadata plus one short probe
// (no network) for users, home folder permissions, tools and startup scripts. The template checks and,
// later, the AI builder work from these instead of guessing from web search.
type imageFacts struct {
	Image        string            `json:"image"`
	User         string            `json:"user"` // image's USER; "" = root
	Env          map[string]string `json:"env"`
	ExposedPorts []string          `json:"exposedPorts"`
	Volumes      []string          `json:"volumes"`
	Entrypoint   []string          `json:"entrypoint"`
	Cmd          []string          `json:"cmd"`
	WorkingDir   string            `json:"workingDir"`
	Users        []imageUser       `json:"users"`   // regular accounts (uid 1000-64999)
	Homes        []homeDir         `json:"homes"`   // /home/* and /root
	Tools        map[string]string `json:"tools"`   // e.g. java -> /usr/bin/java
	Scripts      map[string]string `json:"scripts"` // startup scripts, path -> content (truncated)
	EnvRead      []string          `json:"envRead"` // variables those scripts read
	ProbeError   string            `json:"probeError,omitempty"`
	// What the node needs to run this image: set when the image is built around a user other than 1000:1000
	SuggestedContainerUser string `json:"suggestedContainerUser,omitempty"`
	SuggestionReason       string `json:"suggestionReason,omitempty"`
}

type imageUser struct {
	Name string `json:"name"`
	UID  int    `json:"uid"`
	GID  int    `json:"gid"`
	Home string `json:"home"`
}

type homeDir struct {
	Path string `json:"path"`
	UID  int    `json:"uid"`
	GID  int    `json:"gid"`
	Mode string `json:"mode"` // octal, e.g. "700"
}

const (
	nodeUser       = "1000:1000"
	maxScriptBytes = 60000
	probeMarker    = "@@gbr-probe@@"
)

// Variables every shell or runtime provides; reading them says nothing about the game.
var commonEnv = map[string]bool{
	"HOME": true, "PATH": true, "USER": true, "PWD": true, "SHELL": true, "TERM": true, "HOSTNAME": true,
	"LANG": true, "LC_ALL": true, "UID": true, "GID": true, "IFS": true, "OLDPWD": true, "TMPDIR": true,
	"XDG_RUNTIME_DIR": true, "LD_LIBRARY_PATH": true, "DEBIAN_FRONTEND": true, "BASH_SOURCE": true,
}

var (
	envRef     = regexp.MustCompile(`\$\{?([A-Z][A-Z0-9_]{2,})\b`)
	camelSplit = regexp.MustCompile(`([a-z])([A-Z]+)`)
	absPath    = regexp.MustCompile(`(?:^|[\s'"=;&|(])(/[A-Za-z0-9_./-]+)`)
	// ${VAR:-x} / ${VAR-x}: the command copes when VAR is unset
	defaultedRef = regexp.MustCompile(`\$\{([A-Z][A-Z0-9_]{2,}):?-`)
	// VAR=... or export VAR=... inside the command
	assignedVars = regexp.MustCompile(`(?:^|[\s;&|(])(?:export\s+)?([A-Z][A-Z0-9_]{2,})=`)
)

func uidOf(user string) string {
	uid, _, _ := strings.Cut(user, ":")
	return uid
}

// envName is the environment variable a config key reaches the container as: the node-agent's
// toUpperSnakeCase (RCONPort -> RCONPORT, ServerAdminPassword -> SERVER_ADMIN_PASSWORD).
func envName(key string) string {
	return strings.ToUpper(camelSplit.ReplaceAllString(key, "${1}_${2}"))
}

func runInspectImage(args []string) error {
	fs := flag.NewFlagSet("inspect-image", flag.ExitOnError)
	gamesDir := fs.String("games", "games", "directory holding one folder per game")
	gameID := fs.String("game", "", "also check this game's template against the image")
	asJSON := fs.Bool("json", false, "print the facts as JSON")
	noPull := fs.Bool("no-pull", false, "use the local image instead of pulling it")
	fs.Parse(args)

	var game map[string]any
	image := fs.Arg(0)
	if *gameID != "" {
		g, err := loadMap(filepath.Join(*gamesDir, *gameID, "game.yaml"))
		if err != nil {
			return err
		}
		game = g
		if image == "" {
			image, _ = g["defaultDockerImage"].(string)
		}
	}
	if image == "" {
		return errors.New("inspect-image needs an IMAGE or -game ID")
	}

	if !*noPull {
		if out, err := exec.Command("docker", "pull", "-q", image).CombinedOutput(); err != nil {
			return fmt.Errorf("docker pull %s: %v: %s", image, err, strings.TrimSpace(string(out)))
		}
	}
	facts, err := inspectImage(image, startupScriptPaths(game))
	if err != nil {
		return err
	}

	var problems, notes []string
	if game != nil {
		problems, notes = checkTemplate(game, facts)
	}

	if *asJSON {
		out := map[string]any{"facts": facts}
		if game != nil {
			out["problems"], out["notes"] = problems, notes
		}
		b, _ := json.MarshalIndent(out, "", "  ")
		fmt.Println(string(b))
	} else {
		printFacts(facts)
		if game != nil {
			fmt.Printf("\nTemplate %s against this image:\n", *gameID)
			for _, p := range problems {
				fmt.Println("  PROBLEM  " + p)
			}
			for _, n := range notes {
				fmt.Println("  note     " + n)
			}
			if len(problems) == 0 {
				fmt.Println("  no problems found")
			}
		}
	}
	if len(problems) > 0 {
		return fmt.Errorf("%d problem(s) with %s", len(problems), *gameID)
	}
	return nil
}

func inspectImage(image string, extraScripts []string) (*imageFacts, error) {
	out, err := exec.Command("docker", "image", "inspect", "--format", "{{json .Config}}", image).Output()
	if err != nil {
		return nil, fmt.Errorf("docker image inspect %s: %w", image, err)
	}
	var cfg struct {
		User         string
		Env          []string
		ExposedPorts map[string]any
		Volumes      map[string]any
		Entrypoint   []string
		Cmd          []string
		WorkingDir   string
	}
	if err := json.Unmarshal(out, &cfg); err != nil {
		return nil, fmt.Errorf("parse image config: %w", err)
	}
	f := &imageFacts{
		Image: image, User: cfg.User, Env: map[string]string{}, Entrypoint: cfg.Entrypoint, Cmd: cfg.Cmd,
		WorkingDir: cfg.WorkingDir, ExposedPorts: sortedSet(cfg.ExposedPorts), Volumes: sortedSet(cfg.Volumes),
		Tools: map[string]string{}, Scripts: map[string]string{},
	}
	for _, kv := range cfg.Env {
		if k, v, ok := strings.Cut(kv, "="); ok {
			f.Env[k] = v
		}
	}

	scripts := append(append([]string{}, firstWord(cfg.Entrypoint), firstWord(cfg.Cmd)), extraScripts...)
	probe, err := runProbe(image, scripts)
	if err != nil {
		f.ProbeError = err.Error()
	} else {
		parseProbe(f, probe)
	}
	suggestContainerUser(f)
	return f, nil
}

// runProbe reads users, home folders, tools and scripts from inside the image. --network none: the probe
// never needs the internet, and the image's own entrypoint is bypassed so nothing of the game starts.
func runProbe(image string, scripts []string) (string, error) {
	script := `
m="` + probeMarker + `"
echo "$m passwd"; cat /etc/passwd 2>/dev/null
echo "$m homes"; for d in /home/* /root; do [ -d "$d" ] && stat -c '%n|%u|%g|%a' "$d" 2>/dev/null; done
echo "$m tools"
for t in java steamcmd steamcmd.sh proton wine wine64 bash python3 node dotnet; do p=$(command -v "$t" 2>/dev/null) && echo "$t|$p"; done
for d in /home/*/steamcmd /opt/steamcmd /steamcmd /usr/games/steamcmd; do [ -e "$d" ] && echo "steamcmd-dir|$d"; done
for s in "$@"; do
  [ -n "$s" ] || continue
  p=$(command -v "$s" 2>/dev/null || echo "$s")
  [ -f "$p" ] || continue
  [ "$(head -c 2 "$p")" = "#!" ] || continue
  echo "$m script $p"; head -c ` + strconv.Itoa(maxScriptBytes) + ` "$p"; echo
done
`
	args := append([]string{"run", "--rm", "--network", "none", "--entrypoint", "sh", image, "-c", script, "probe"}, scripts...)
	cmd := exec.Command("docker", args...)
	var stdout, stderr bytes.Buffer
	cmd.Stdout, cmd.Stderr = &stdout, &stderr
	if err := cmd.Run(); err != nil && stdout.Len() == 0 {
		return "", fmt.Errorf("probe could not run a shell in the image: %v %s", err, strings.TrimSpace(stderr.String()))
	}
	return stdout.String(), nil
}

func parseProbe(f *imageFacts, out string) {
	section, scriptPath := "", ""
	var script strings.Builder
	flush := func() {
		if scriptPath != "" {
			f.Scripts[scriptPath] = script.String()
		}
		scriptPath = ""
		script.Reset()
	}
	for _, line := range strings.Split(out, "\n") {
		if rest, ok := strings.CutPrefix(line, probeMarker+" "); ok {
			flush()
			section = rest
			if p, ok := strings.CutPrefix(rest, "script "); ok {
				section, scriptPath = "script", p
			}
			continue
		}
		switch section {
		case "passwd":
			parts := strings.Split(line, ":")
			if len(parts) < 6 {
				continue
			}
			uid, _ := strconv.Atoi(parts[2])
			gid, _ := strconv.Atoi(parts[3])
			if uid >= 1000 && uid < 65000 {
				f.Users = append(f.Users, imageUser{Name: parts[0], UID: uid, GID: gid, Home: parts[5]})
			}
		case "homes":
			parts := strings.Split(line, "|")
			if len(parts) == 4 {
				uid, _ := strconv.Atoi(parts[1])
				gid, _ := strconv.Atoi(parts[2])
				f.Homes = append(f.Homes, homeDir{Path: parts[0], UID: uid, GID: gid, Mode: parts[3]})
			}
		case "tools":
			if name, path, ok := strings.Cut(line, "|"); ok {
				f.Tools[name] = path
			}
		case "script":
			script.WriteString(line + "\n")
		}
	}
	flush()

	seen := map[string]bool{}
	for _, body := range f.Scripts {
		for _, m := range envRef.FindAllStringSubmatch(body, -1) {
			if !commonEnv[m[1]] {
				seen[m[1]] = true
			}
		}
	}
	f.EnvRead = sortedKeysBool(seen)
}

// suggestContainerUser works out whether the image only works as a specific user. Images either declare
// it (USER) or are built around an account whose home folder nobody else may enter (e.g. a home folder
// at mode 700 owned by a uid other than 1000).
func suggestContainerUser(f *imageFacts) {
	if u := strings.TrimSpace(f.User); u != "" && u != "root" && u != "0" {
		resolved := resolveUser(f, u)
		if resolved != "" && resolved != nodeUser {
			f.SuggestedContainerUser = resolved
			f.SuggestionReason = fmt.Sprintf("the image declares USER %s", u)
		}
		return
	}
	for _, h := range f.Homes {
		if h.Path == "/root" || h.UID == 0 || strconv.Itoa(h.UID) == uidOf(nodeUser) {
			continue
		}
		mode, err := strconv.ParseInt(h.Mode, 8, 32)
		if err != nil || mode&0o005 != 0 { // others can enter it, so 1000:1000 would get in too
			continue
		}
		// The account's own group if it has one; the folder's group otherwise
		gid := h.GID
		for _, acc := range f.Users {
			if acc.UID == h.UID {
				gid = acc.GID
			}
		}
		f.SuggestedContainerUser = fmt.Sprintf("%d:%d", h.UID, gid)
		f.SuggestionReason = fmt.Sprintf("%s belongs to uid %d with mode %s, so the node's default %s cannot enter it", h.Path, h.UID, h.Mode, nodeUser)
		return
	}
}

func resolveUser(f *imageFacts, u string) string {
	name, group, _ := strings.Cut(u, ":")
	for _, acc := range f.Users {
		if acc.Name == name || strconv.Itoa(acc.UID) == name {
			if group == "" {
				return fmt.Sprintf("%d:%d", acc.UID, acc.GID)
			}
			if _, err := strconv.Atoi(group); err == nil {
				return fmt.Sprintf("%d:%s", acc.UID, group)
			}
			return fmt.Sprintf("%d:%d", acc.UID, acc.GID)
		}
	}
	if _, err := strconv.Atoi(name); err == nil {
		if group == "" {
			group = name
		}
		return name + ":" + group
	}
	return ""
}

// checkTemplate compares a game template with what the image needs. Problems are things that stop or
// silently break the server; notes are worth a look but often fine.
func checkTemplate(game map[string]any, f *imageFacts) (problems, notes []string) {
	set, _ := game["containerUser"].(string)
	effective := set
	if effective == "" {
		effective = nodeUser
	}
	// Only the uid decides whether the server can enter its folders; the group may reasonably differ
	if f.SuggestedContainerUser != "" && uidOf(effective) != uidOf(f.SuggestedContainerUser) {
		problems = append(problems, fmt.Sprintf("image needs containerUser %q (%s), template runs as %s",
			f.SuggestedContainerUser, f.SuggestionReason, effective))
	}
	if f.ProbeError != "" {
		notes = append(notes, "could not look inside the image: "+f.ProbeError)
	}

	// Everything the container will have in its environment: the image's own defaults plus one variable
	// per config key, named the way the node-agent names them
	provided := map[string]bool{}
	for k := range f.Env {
		provided[k] = true
	}
	defaults, _ := game["defaultConfig"].(map[string]any)
	byLoose := map[string]string{}
	for key := range defaults {
		n := envName(key)
		provided[n] = true
		byLoose[strings.ReplaceAll(n, "_", "")] = n
	}
	if cmd, _ := game["defaultStartupCommand"].(string); cmd != "" {
		for _, m := range assignedVars.FindAllStringSubmatch(cmd, -1) {
			provided[m[1]] = true // set by the command itself, e.g. export START_PARAMS=...
		}
		withDefault := map[string]bool{}
		for _, m := range defaultedRef.FindAllStringSubmatch(cmd, -1) {
			withDefault[m[1]] = true
		}
		for _, v := range uniqueRefs(cmd) {
			if provided[v] || commonEnv[v] {
				continue
			}
			near, misnamed := byLoose[strings.ReplaceAll(v, "_", "")]
			switch {
			case misnamed:
				// The RCON bug: ${RCON_PORT:-27020} while the RCONPort setting arrives as RCONPORT, so the
				// default always wins and the owner's setting is silently ignored
				problems = append(problems, fmt.Sprintf("startup command reads ${%s}, but the setting arrives as %s", v, near))
			case !withDefault[v]:
				problems = append(problems, fmt.Sprintf("startup command reads ${%s}, which nothing sets and has no default", v))
			}
		}
	}

	ports, _ := game["defaultPorts"].([]any)
	templatePorts := map[string]bool{}
	for _, p := range ports {
		templatePorts[containerPort(fmt.Sprint(p))] = true
	}
	for _, p := range f.ExposedPorts {
		if !templatePorts[p] {
			notes = append(notes, fmt.Sprintf("image exposes %s, template does not publish it (fine for RCON/admin ports)", p))
		}
	}
	return problems, notes
}

// containerPort turns "8211:8211/udp" or "7777/udp" or "25565" into Docker's "8211/udp" form.
func containerPort(p string) string {
	if i := strings.LastIndex(p, ":"); i >= 0 {
		p = p[i+1:]
	}
	if !strings.Contains(p, "/") {
		p += "/tcp"
	}
	return p
}

// startupScriptPaths picks absolute paths out of a template's startup command (e.g. an image's
// /usr/bin/start_server) so the probe reads those scripts too.
func startupScriptPaths(game map[string]any) []string {
	cmd, _ := game["defaultStartupCommand"].(string)
	var paths []string
	for _, m := range absPath.FindAllStringSubmatch(cmd, -1) {
		if p := m[1]; !strings.HasPrefix(p, "/server") && !strings.HasPrefix(p, "/data") && !strings.HasPrefix(p, "/home/") {
			paths = append(paths, p)
		}
	}
	return paths
}

func uniqueRefs(s string) []string {
	seen := map[string]bool{}
	for _, m := range envRef.FindAllStringSubmatch(s, -1) {
		seen[m[1]] = true
	}
	return sortedKeysBool(seen)
}

func printFacts(f *imageFacts) {
	fmt.Printf("Image:        %s\n", f.Image)
	user := f.User
	if user == "" {
		user = "(root)"
	}
	fmt.Printf("Default user: %s\n", user)
	if f.SuggestedContainerUser != "" {
		fmt.Printf("Needs:        containerUser %q, because %s\n", f.SuggestedContainerUser, f.SuggestionReason)
	}
	fmt.Printf("Entrypoint:   %s\n", strings.Join(append(append([]string{}, f.Entrypoint...), f.Cmd...), " "))
	fmt.Printf("Ports:        %s\n", strings.Join(f.ExposedPorts, ", "))
	fmt.Printf("Volumes:      %s\n", strings.Join(f.Volumes, ", "))
	for _, u := range f.Users {
		fmt.Printf("User:         %s (uid %d, gid %d, home %s)\n", u.Name, u.UID, u.GID, u.Home)
	}
	for _, h := range f.Homes {
		fmt.Printf("Home:         %s owned by %d:%d mode %s\n", h.Path, h.UID, h.GID, h.Mode)
	}
	if len(f.Tools) > 0 {
		var t []string
		for _, k := range sortedKeysString(f.Tools) {
			t = append(t, k+"="+f.Tools[k])
		}
		fmt.Printf("Tools:        %s\n", strings.Join(t, ", "))
	}
	for _, p := range sortedKeysString(f.Scripts) {
		fmt.Printf("Script:       %s (%d bytes)\n", p, len(f.Scripts[p]))
	}
	if len(f.EnvRead) > 0 {
		fmt.Printf("Scripts read: %s\n", strings.Join(f.EnvRead, ", "))
	}
	if len(f.Env) > 0 {
		fmt.Printf("Image env:    %s\n", strings.Join(sortedKeysString(f.Env), ", "))
	}
	if f.ProbeError != "" {
		fmt.Printf("Probe:        %s\n", f.ProbeError)
	}
}

func firstWord(s []string) string {
	if len(s) == 0 {
		return ""
	}
	return s[0]
}

func sortedSet(m map[string]any) []string {
	out := make([]string, 0, len(m))
	for k := range m {
		out = append(out, k)
	}
	sort.Strings(out)
	return out
}

func sortedKeysBool(m map[string]bool) []string {
	out := make([]string, 0, len(m))
	for k := range m {
		out = append(out, k)
	}
	sort.Strings(out)
	return out
}

func sortedKeysString(m map[string]string) []string {
	out := make([]string, 0, len(m))
	for k := range m {
		out = append(out, k)
	}
	sort.Strings(out)
	return out
}