package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"time"
)

// runCmd runs a command and returns its stdout (trimmed by callers). Used to read a
// built binary's `version` output.
func runCmd(name string, args ...string) (string, error) {
	out, err := exec.Command(name, args...).Output()
	if err != nil {
		return "", fmt.Errorf("%s %s: %v", name, strings.Join(args, " "), err)
	}
	return string(out), nil
}

// Versioned, side-by-side deploys.
//
// The box runs one `lchat` binary per instance, but they need not be the SAME
// binary: each instance's systemd unit execs /opt/lchat/live/<host>, a symlink
// into a content store at /opt/lchat/versions/<ver>/lchat. Swapping one link (and
// restarting one unit) moves one instance to a new build while every other
// instance stays put — so a new build can be canaried on `test` before it's
// promoted fleet-wide, and rolled back per instance. The registry records each
// instance's current + previous version (source of truth for `lchat instances`
// and `rollback`). See docs/instance-router.md § Versioned deploys.

// buildVersion is the binary's build id, compiled in via
// -ldflags "-X main.buildVersion=$(git rev-parse --short HEAD)" (see
// deploy/build-linux.sh). "dev" means an un-stamped local build. `lchat version`
// prints it; the running instance reports it at /api/version.
var buildVersion = "dev"

// optDir is the on-box install root: the version store and the per-instance
// "live" symlinks live under it. A var with an env override so tests can point it
// at a temp dir instead of the real /opt/lchat.
var optDir = envOr("LCHAT_OPT_DIR", "/opt/lchat")

func versionsDir() string              { return filepath.Join(optDir, "versions") }
func liveDir() string                  { return filepath.Join(optDir, "live") }
func versionBinPath(ver string) string { return filepath.Join(versionsDir(), ver, "lchat") }

// liveLink is the per-instance pointer the systemd unit execs
// (/opt/lchat/live/<host>) — a symlink to a versionBinPath. Swapping it atomically
// is how one instance moves between versions independently of the others.
func liveLink(host string) string { return filepath.Join(liveDir(), host) }

// versionRe keeps a version usable as a directory name (no slashes, no traversal).
var versionRe = regexp.MustCompile(`^[a-zA-Z0-9._-]+$`)

func validVersion(v string) bool {
	return v != "" && len(v) <= 64 && v != "." && v != ".." && versionRe.MatchString(v)
}

// readBinaryVersion runs `<path> version` to read a built binary's compiled-in id,
// so `stage` can key the store by it without the operator retyping the sha.
func readBinaryVersion(path string) (string, error) {
	out, err := runCmd(path, "version")
	if err != nil {
		return "", err
	}
	return strings.TrimSpace(out), nil
}

// copyFileAtomic copies src to dst via a temp file + rename, so a reader never sees
// a half-written binary (and a crashed copy leaves the old one intact).
func copyFileAtomic(src, dst string, mode os.FileMode) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	tmp := dst + ".tmp"
	out, err := os.OpenFile(tmp, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, mode)
	if err != nil {
		return err
	}
	if _, err := io.Copy(out, in); err != nil {
		out.Close()
		os.Remove(tmp)
		return err
	}
	if err := out.Close(); err != nil {
		os.Remove(tmp)
		return err
	}
	// O_CREATE's mode is masked by umask; set it explicitly so the binary is +x.
	if err := os.Chmod(tmp, mode); err != nil {
		os.Remove(tmp)
		return err
	}
	if err := os.Rename(tmp, dst); err != nil {
		os.Remove(tmp)
		return err
	}
	return nil
}

// stageVersion copies a built binary into the store under ver, so instances can be
// pointed at it. Idempotent: re-staging the same ver overwrites its binary.
func stageVersion(srcPath, ver string) error {
	if !validVersion(ver) {
		return fmt.Errorf("invalid version %q (allowed: letters, digits, . _ -)", ver)
	}
	dir := filepath.Join(versionsDir(), ver)
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return err
	}
	// MkdirAll's mode is masked by the caller's umask, so a tight operator
	// umask (e.g. 077 from a hardened shell) leaves the version dir 0700 —
	// the sandboxed instance user then can't traverse in to exec its binary
	// (203/EXEC "Permission denied", found staging a canary). Force it, like
	// copyFileAtomic already forces the binary's +x.
	if err := os.Chmod(dir, 0o755); err != nil {
		return err
	}
	return copyFileAtomic(srcPath, versionBinPath(ver), 0o755)
}

// pointLive atomically repoints an instance's live symlink at ver's binary. It
// verifies ver is staged first, so a typo can't leave a dangling link that would
// fail the unit's ExecStart on the next restart.
func pointLive(host, ver string) error {
	target := versionBinPath(ver)
	if _, err := os.Stat(target); err != nil {
		return fmt.Errorf("version %q is not staged (%s): run `lchat stage --file <binary>` first", ver, target)
	}
	if err := os.MkdirAll(liveDir(), 0o755); err != nil {
		return err
	}
	link := liveLink(host)
	tmp := link + ".tmp"
	_ = os.Remove(tmp) // a leftover temp link from a crashed swap
	if err := os.Symlink(target, tmp); err != nil {
		return err
	}
	if err := os.Rename(tmp, link); err != nil {
		os.Remove(tmp)
		return err
	}
	return nil
}

// liveVersion reports the version an instance's live symlink currently points at,
// or "" (no error) if the instance hasn't been put under version management yet.
func liveVersion(host string) (string, error) {
	target, err := os.Readlink(liveLink(host))
	if err != nil {
		if os.IsNotExist(err) {
			return "", nil
		}
		return "", err
	}
	// target = /opt/lchat/versions/<ver>/lchat → the parent dir's name is <ver>.
	return filepath.Base(filepath.Dir(target)), nil
}

// selfPath is the running control binary — what add-instance / init-versions stage
// so a fresh instance runs the same build as the operator's `lchat`. A var for tests.
var selfPath = func() (string, error) { return os.Executable() }

// ensureSelfStaged stages the running control binary under its build id and returns
// that id, so add-instance can point a fresh instance at a real staged version
// before the unit (ExecStart=/opt/lchat/live/<host>) first starts.
func ensureSelfStaged() (string, error) {
	self, err := selfPath()
	if err != nil {
		return "", err
	}
	if err := stageVersion(self, buildVersion); err != nil {
		return "", err
	}
	return buildVersion, nil
}

// runStage: `lchat stage --file <binary>` — register a built binary in the store
// under its build id (or --version), ready for `lchat deploy --version`.
func runStage(args []string) {
	fs := flag.NewFlagSet("stage", flag.ExitOnError)
	file := fs.String("file", "", "path to the built lchat binary to register (required)")
	verFlag := fs.String("version", "", "version id to store it under (default: the binary's own `lchat version`)")
	_ = fs.Parse(reorderFlagsFirst(fs, args))

	if *file == "" {
		die("stage: --file <path> is required")
	}
	ver := strings.TrimSpace(*verFlag)
	if ver == "" {
		v, err := readBinaryVersion(*file)
		if err != nil {
			die("stage: can't read the build id from %s: %v\n  build it with -ldflags \"-X main.buildVersion=$(git rev-parse --short HEAD)\" (see deploy/build-linux.sh), or pass --version", *file, err)
		}
		ver = v
	}
	if ver == "" || ver == "dev" {
		die("stage: refusing to stage build id %q — build with -ldflags \"-X main.buildVersion=$(git rev-parse --short HEAD)\" or pass --version <id>", ver)
	}
	if !validVersion(ver) {
		die("stage: invalid version %q (allowed: letters, digits, . _ -)", ver)
	}
	if err := stageVersion(*file, ver); err != nil {
		die("stage: %v", err)
	}
	fmt.Printf("staged %s -> %s\n", ver, versionBinPath(ver))
}

// runInitVersions: `lchat init-versions` — one-time adoption. Stage the running
// control binary and point every instance's live symlink at it, so the versioned
// unit (ExecStart=/opt/lchat/live/<host>) resolves for instances that predate it.
// Idempotent: it only creates a live link that's missing, never repoints a managed
// instance (use `deploy --version` for that).
func runInitVersions(args []string) {
	fs := flag.NewFlagSet("init-versions", flag.ExitOnError)
	registryPath := fs.String("registry", envOr("LCHAT_REGISTRY", "/data/registry.db"), "router registry SQLite path")
	from := fs.String("from", "", "binary to seed the store with (default: the running control binary)")
	verFlag := fs.String("version", "", "version id for the seed binary (default: its own build id)")
	_ = fs.Parse(reorderFlagsFirst(fs, args))

	src := strings.TrimSpace(*from)
	ver := strings.TrimSpace(*verFlag)
	if src == "" {
		self, err := selfPath()
		if err != nil {
			die("init-versions: locate self: %v", err)
		}
		src = self
		if ver == "" {
			ver = buildVersion
		}
	} else if ver == "" {
		v, err := readBinaryVersion(src)
		if err != nil {
			die("init-versions: read build id from %s: %v (or pass --version)", src, err)
		}
		ver = v
	}
	if !validVersion(ver) {
		die("init-versions: invalid seed version %q (build with -ldflags, or pass --version)", ver)
	}
	if err := stageVersion(src, ver); err != nil {
		die("init-versions: stage: %v", err)
	}
	fmt.Printf("seeded %s -> %s\n", ver, versionBinPath(ver))

	reg, err := openRegistry(*registryPath)
	if err != nil {
		die("init-versions: open registry: %v", err)
	}
	routes, err := reg.list()
	if err != nil {
		die("init-versions: %v", err)
	}
	for _, rt := range routes {
		if rt.Kind != "instance" {
			continue
		}
		cur, err := liveVersion(rt.Host)
		if err != nil {
			die("init-versions: %s: read live link: %v", rt.Host, err)
		}
		if cur != "" {
			fmt.Printf("  %s already on %s (left as-is)\n", rt.Host, cur)
			continue
		}
		if err := pointLive(rt.Host, ver); err != nil {
			die("init-versions: %s: %v", rt.Host, err)
		}
		rt.Version = ver
		if err := reg.upsert(rt); err != nil {
			die("init-versions: %s: record version: %v", rt.Host, err)
		}
		fmt.Printf("  %s -> %s\n", rt.Host, ver)
	}
	fmt.Println("init-versions: done (restart instances onto the versioned unit with `lchat deploy --all`)")
}

// runRollback: `lchat rollback <host>` (or --all) — repoint the instance(s) at the
// previous version recorded at the last `deploy --version`, restart, health-check.
// It swaps current<->previous so a second rollback returns to where you were.
func runRollback(args []string) {
	fs := flag.NewFlagSet("rollback", flag.ExitOnError)
	registryPath := fs.String("registry", envOr("LCHAT_REGISTRY", "/data/registry.db"), "router registry SQLite path")
	unitTmpl := fs.String("unit", envOr("LCHAT_UNIT", "lchat@"), "systemd instance unit template prefix")
	all := fs.Bool("all", false, "roll back every instance that has a previous version")
	_ = fs.Parse(reorderFlagsFirst(fs, args))

	unlock, err := acquireDeployLock()
	if err != nil {
		die("rollback: %v", err)
	}
	defer unlock()

	reg, err := openRegistry(*registryPath)
	if err != nil {
		die("rollback: open registry: %v", err)
	}
	hosts, err := deployTargets(reg, *all, fs.Arg(0), "rollback")
	if err != nil {
		die("%v", err)
	}
	for _, host := range hosts {
		rt, err := reg.get(host)
		if err != nil || rt == nil {
			die("rollback %s: unknown host (run add-instance first)", host)
		}
		if rt.Kind != "instance" {
			die("rollback %s: not an instance (kind %q)", host, rt.Kind)
		}
		prev := strings.TrimSpace(rt.PrevVersion)
		if prev == "" {
			if *all {
				continue // nothing to roll back to; skip in a fleet rollback
			}
			die("rollback %s: no previous version recorded", host)
		}
		cur, _ := liveVersion(host)
		if rt.Hibernate {
			// No unit to restart — a pure repoint, mirroring the auto-updater's
			// hibernating phase. This is the rescue path when a repointed build
			// turns out bad: the next wake execs the previous version.
			if err := pointLive(host, prev); err != nil {
				die("rollback %s: %v", host, err)
			}
			if err := reg.setVersions(host, prev, cur); err != nil {
				die("rollback %s: record: %v", host, err)
			}
			fmt.Printf("rolled back %s to %s (hibernating — runs it on next wake)\n", host, prev)
			continue
		}
		if err := pointLive(host, prev); err != nil {
			die("rollback %s: %v", host, err)
		}
		if err := systemctl("restart", *unitTmpl+host); err != nil {
			die("rollback %s: restart: %v", host, err)
		}
		if err := waitHealthy(rt.Socket, 20*time.Second); err != nil {
			die("rollback %s: %v", host, err)
		}
		if err := reg.setVersions(host, prev, cur); err != nil {
			die("rollback %s: record: %v", host, err)
		}
		restartBotSidecar(host)
		fmt.Printf("rolled back %s to %s\n", host, prev)
	}
}

// deployTargets resolves the host list shared by deploy/rollback: every instance
// (--all) or the one validated positional host.
func deployTargets(reg *Registry, all bool, arg, cmd string) ([]string, error) {
	if all {
		routes, err := reg.list()
		if err != nil {
			return nil, fmt.Errorf("%s: %v", cmd, err)
		}
		var hosts []string
		for _, r := range routes {
			if r.Kind == "instance" {
				hosts = append(hosts, r.Host)
			}
		}
		if len(hosts) == 0 {
			return nil, fmt.Errorf("%s: no instances", cmd)
		}
		return hosts, nil
	}
	h := strings.ToLower(strings.TrimSpace(arg))
	if !validHost(h) {
		return nil, fmt.Errorf("%s: invalid host %q (or pass --all)", cmd, h)
	}
	return []string{h}, nil
}

// waitHealthyVersion is waitHealthy plus a build check: it waits until the instance
// not only answers but reports the expected build at /api/version — so a deploy
// fails loudly if the process is still the old binary (a stale socket, or a restart
// that didn't take), rather than silently "succeeding".
func waitHealthyVersion(socket, wantBuild string, timeout time.Duration) error {
	deadline := time.Now().Add(timeout)
	client := unixHTTPClient(socket)
	var last error
	for time.Now().Before(deadline) {
		resp, err := client.Get("http://lchat/api/version")
		if err == nil {
			var body struct {
				Build string `json:"build"`
			}
			derr := json.NewDecoder(io.LimitReader(resp.Body, 1<<16)).Decode(&body)
			resp.Body.Close()
			if derr == nil && body.Build == wantBuild {
				return nil
			}
			if derr != nil {
				last = fmt.Errorf("read /api/version: %v", derr)
			} else {
				last = fmt.Errorf("live build %q != expected %q", body.Build, wantBuild)
			}
		} else {
			last = err
		}
		time.Sleep(150 * time.Millisecond)
	}
	return fmt.Errorf("instance didn't report build %q within %s: %v", wantBuild, timeout, last)
}

// probeVersion reads the build the process behind a socket is ACTUALLY serving
// (its /api/version), not the symlink target. This is how `update status
// --probe` catches an awake hibernating instance that repointed its symlink but
// still runs the old binary (docs/awake-instance-convergence.md). "" + ok=false
// when the socket doesn't answer (asleep, or no such instance).
func probeVersion(socket string) (build string, ok bool) {
	resp, err := unixHTTPClient(socket).Get("http://lchat/api/version")
	if err != nil {
		return "", false
	}
	defer resp.Body.Close()
	var body struct {
		Build string `json:"build"`
	}
	if json.NewDecoder(io.LimitReader(resp.Body, 1<<16)).Decode(&body) != nil {
		return "", false
	}
	return body.Build, body.Build != ""
}

// unixHTTPClient dials a unix socket over HTTP — the router↔instance transport,
// reused by the health checks.
func unixHTTPClient(socket string) *http.Client {
	return &http.Client{
		Timeout: 3 * time.Second,
		Transport: &http.Transport{
			DialContext: func(ctx context.Context, _, _ string) (net.Conn, error) {
				return (&net.Dialer{}).DialContext(ctx, "unix", socket)
			},
		},
	}
}

// ---- sibling installs -------------------------------------------------------

// siblingBinaries names the box's non-instance service binaries that `lchat
// install` may replace, each with the unit that runs it. These are the fixed
// ExecStart targets of their units — the door the SSH hardening deliberately
// closed for ad-hoc root shells, reopened as one audited, allowlisted verb
// (operators run it via `sudo -n /opt/lchat/lchat install …`). Extending the
// set is a reviewed PR, never a box edit.
var siblingBinaries = map[string]string{
	"home":    "lchat-home",
	"push":    "lchat-push",
	"connect": "lchat-connect",
}

// installSibling swaps one sibling service's binary to a staged version:
// back up the current binary (cp to <name>.prev-<date>, the convention the
// box already uses), then atomically copy the staged build over the unit's
// ExecStart path. It deliberately does NOT restart the unit — the operator
// does that visibly via lchat-svc, so an install can be staged ahead of a
// maintenance window. Returns the human summary to print.
func installSibling(name, ver string) (string, error) {
	unit, ok := siblingBinaries[name]
	if !ok {
		return "", fmt.Errorf("unknown sibling %q (know: home, push, connect)", name)
	}
	src := versionBinPath(ver)
	if _, err := os.Stat(src); err != nil {
		return "", fmt.Errorf("version %q is not staged — run `lchat stage --file <binary>` first", ver)
	}
	target := filepath.Join(optDir, name)
	prev := ""
	if _, err := os.Stat(target); err == nil {
		prev = target + ".prev-" + time.Now().Format("20060102-150405")
		if err := copyFileAtomic(target, prev, 0o755); err != nil {
			return "", fmt.Errorf("back up current binary: %v", err)
		}
	}
	if err := copyFileAtomic(src, target, 0o755); err != nil {
		return "", fmt.Errorf("install: %v", err)
	}
	msg := fmt.Sprintf("installed %s %s -> %s", name, ver, target)
	if prev != "" {
		msg += fmt.Sprintf(" (previous at %s)", prev)
	}
	msg += fmt.Sprintf("\nrestart to pick it up: sudo -n /usr/local/sbin/lchat-svc restart %s", unit)
	return msg, nil
}

// runInstall: `lchat install <name> --version <ver>` — put a staged build onto
// a sibling service's fixed binary path. The staged version store doesn't care
// what a binary IS (stage --file records any build under its own version), so
// sibling binaries ride the same store the messenger versions do.
func runInstall(args []string) {
	fs := flag.NewFlagSet("install", flag.ExitOnError)
	verFlag := fs.String("version", "", "staged version to install (see `lchat stage`)")
	_ = fs.Parse(reorderFlagsFirst(fs, args))
	name := strings.TrimSpace(fs.Arg(0))
	ver := strings.TrimSpace(*verFlag)
	if name == "" || ver == "" {
		die("install: usage: lchat install <home|push|connect> --version <ver>")
	}
	if !validVersion(ver) {
		die("install: invalid version %q", ver)
	}
	msg, err := installSibling(name, ver)
	if err != nil {
		die("install: %v", err)
	}
	fmt.Println(msg)
}