package main

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	"wxch.org/wxch/internal/modules"
)

// Engines run directly on the machine (no Docker): the agent installs a Python environment per engine with uv, starts the
// engine as its own child process (own process group, own environment without the agent's token) and talks to it on
// 127.0.0.1 exactly as it does to the container. hostRunner is a modRunner, so the rest of the module code (jobs, warm-up,
// idle stop, the comfy. tunnel) does not know which kind of engine it is talking to.
//
// Everything lives under the modules folder:  .host/{uv,python,cache}  and  <engine>/{venv,home,tmp,engine.log,install.log}.

const (
	uvVersion   = "0.12.23"
	hostReady   = ".wxch-ready"
	hostLogMax  = 20 << 20
	hostTorchBk = "cu130" // PyTorch wheels built for CUDA 13 (the GB10 / Blackwell driver stack)
)

// uvSums pins the uv release by checksum: nothing downloaded from the internet is run without it matching.
var uvSums = map[string]string{
	"arm64": "6524bd338177ed50d035d39354e12545e993bbeba2ecbddf0480c5b3a81d313f",
	"amd64": "9167d72b3319674b6303c4cbe071854bba13ebdf3d76b1a7cbdc175471fb66d6",
}
var uvTriples = map[string]string{"arm64": "aarch64", "amd64": "x86_64"}

// uvURL is a variable for tests.
var uvURL = func(arch string) string {
	return "https://github.com/astral-sh/uv/releases/download/" + uvVersion + "/uv-" + uvTriples[arch] + "-unknown-linux-gnu.tar.gz"
}

type hostProc struct {
	cmd     *exec.Cmd
	pgid    int
	done    chan struct{}
	started time.Time
}

type hostInstall struct {
	active bool
	step   string
	err    string
}

var host = struct {
	sync.Mutex
	procs    map[string]*hostProc
	installs map[string]*hostInstall
}{procs: map[string]*hostProc{}, installs: map[string]*hostInstall{}}

type hostRunner struct {
	cfg nodeConfig
	cat *modules.Catalog
}

type hostPaths struct{ root, hostDir, dir, venv, log, install, home, tmp, pid string }

func (h hostRunner) paths(engine string) hostPaths {
	root := modulesDir(h.cfg)
	dir := filepath.Join(root, engine)
	return hostPaths{root: root, hostDir: filepath.Join(root, ".host"), dir: dir, venv: filepath.Join(dir, "venv"),
		log: filepath.Join(dir, "engine.log"), install: filepath.Join(dir, "install.log"),
		home: filepath.Join(dir, "home"), tmp: filepath.Join(dir, "tmp"), pid: filepath.Join(dir, "engine.pid")}
}

func hostPython(venv string) string { return filepath.Join(venv, "bin", "python") }

func (h hostRunner) ready(engine string) bool {
	_, err := os.Stat(filepath.Join(h.paths(engine).venv, hostReady))
	return err == nil
}

// engineEnv is the whole environment of an engine. It is built from nothing: the agent's own environment holds the
// machine's token, which an engine (and the workflows people run on it) must never see.
func engineEnv(p hostPaths) []string {
	return []string{
		"PATH=" + filepath.Join(p.venv, "bin") + ":/usr/local/bin:/usr/bin:/bin",
		"HOME=" + p.home, "TMPDIR=" + p.tmp, "LANG=C.UTF-8", "PYTHONUNBUFFERED=1", "PYTHONNOUSERSITE=1",
		"DO_NOT_TRACK=1", "HF_HUB_OFFLINE=1", "VIRTUAL_ENV=" + p.venv,
	}
}

// installEnv is for uv and pip: the same, plus where uv keeps Python and its cache (inside the modules folder, so that a
// machine with an external models disk keeps all of it there and the cache and the environments share one disk).
func installEnv(p hostPaths) []string {
	return []string{"PATH=/usr/local/bin:/usr/bin:/bin", "HOME=" + p.home, "TMPDIR=" + p.tmp, "LANG=C.UTF-8",
		"UV_PYTHON_INSTALL_DIR=" + filepath.Join(p.hostDir, "python"), "UV_CACHE_DIR=" + filepath.Join(p.hostDir, "cache"),
		"UV_PYTHON_PREFERENCE=only-managed", "UV_NO_PROGRESS=1", "UV_NO_CONFIG=1", "NO_COLOR=1", "DO_NOT_TRACK=1"}
}

// comfyTuning are the only extra flags the owner may put in <engine folder>/engine-args.txt (one or more, separated by
// spaces or lines): memory and speed switches of ComfyUI, to try on a machine without a new release. Nothing that changes
// where ComfyUI listens or what it loads can be added this way.
var comfyTuning = map[string]bool{
	"--disable-dynamic-vram": true, "--enable-dynamic-vram": true, "--disable-fast-disk": true, "--fast-disk": true,
	"--disable-pinned-memory": true, "--disable-mmap": true, "--disable-smart-memory": true, "--disable-async-offload": true,
	"--disable-cuda-graphs": true, "--disable-comfy-compiler": true, "--disable-cuda-malloc": true, "--disable-triton-backend": true,
	"--highvram": true, "--gpu-only": true,
}

func comfyExtraArgs(dir string) []string {
	b, err := os.ReadFile(filepath.Join(dir, "engine-args.txt"))
	if err != nil {
		return nil
	}
	var out []string
	for _, a := range strings.Fields(string(b)) {
		if comfyTuning[a] {
			out = append(out, a)
		} else {
			logf("engine-args.txt: ignoring %q (not an allowed switch)", a)
		}
	}
	return out
}

// engineArgv is what is run for an engine (a variable for tests). The listen address is always 127.0.0.1.
var engineArgv = func(cat *modules.Catalog, e modules.Engine, p hostPaths) (argv []string, dir string) {
	py, port := hostPython(p.venv), strconv.Itoa(cat.HostPort(e.ID))
	switch e.ID {
	case "comfyui":
		argv := []string{py, "main.py", "--listen", "127.0.0.1", "--port", port, "--disable-auto-launch", "--disable-api-nodes", "--disable-metadata"}
		return append(argv, comfyExtraArgs(p.dir)...), filepath.Join(p.dir, filepath.Base(e.Workdir))
	case "whisper":
		return []string{py, "whisper_server.py", "--host", "127.0.0.1", "--port", port, "--model", filepath.Join(p.dir, "models", "whisper-large-v3-turbo")}, p.dir
	}
	return nil, ""
}

// ---- modRunner ----

func (h hostRunner) Status() ([]runnerModule, error) {
	host.Lock()
	defer host.Unlock()
	var out []runnerModule
	for _, e := range h.cat.Engines {
		m := runnerModule{Engine: e.ID, Port: h.cat.HostPort(e.ID), ImagePresent: h.ready(e.ID)}
		if pr := host.procs[e.ID]; pr != nil {
			select {
			case <-pr.done:
				m.State = "exited"
			default:
				m.State = "running"
			}
		}
		if in := host.installs[e.ID]; in != nil {
			m.Pulling, m.PullErr, m.PullStep = in.active, in.err, in.step
		}
		out = append(out, m)
	}
	return out, nil
}

func (h hostRunner) Pull(engine string) error {
	e, ok := h.cat.Engine(engine)
	if !ok {
		return fmt.Errorf("unknown engine")
	}
	h.kick(e)
	return nil
}

var provisionMu sync.Mutex

// hostProvisionAtStart installs the runtime of every engine that has a module installed but no environment yet (a machine
// that ran its engines in Docker before, or a first install that was interrupted), so that nobody has to be the first to ask.
func hostProvisionAtStart() {
	time.Sleep(20 * time.Second)
	cfg := loadNodeConfig()
	if engineMode(cfg) != "host" {
		return
	}
	cat, err := loadCatalog()
	if err != nil {
		return
	}
	st := readModuleState(cfg)
	h := hostRunner{cfg: cfg, cat: cat}
	for _, e := range cat.Engines {
		if len(capabilitiesUsing(cat, st, e.ID)) > 0 && !h.ready(e.ID) {
			logf("engine %s: installing its runtime (Python and PyTorch) in the background", e.ID)
			h.kick(e)
		}
	}
}

// kick installs the engine's runtime in the background unless that is already going on.
func (h hostRunner) kick(e modules.Engine) {
	host.Lock()
	in := host.installs[e.ID]
	if in != nil && in.active {
		host.Unlock()
		return
	}
	in = &hostInstall{active: true, step: "starting"}
	host.installs[e.ID] = in
	host.Unlock()
	go func() {
		ctx, cancel := context.WithTimeout(context.Background(), 3*time.Hour)
		defer cancel()
		provisionMu.Lock() // one install at a time: they share uv, its cache and the downloads
		err := h.provision(ctx, e, func(s string) { host.Lock(); in.step = s; host.Unlock() })
		provisionMu.Unlock()
		if err == nil {
			h.stopLegacyContainer(e.ID) // the old container still holds the port and would keep answering
		}
		host.Lock()
		in.active = false
		if err != nil {
			in.err = err.Error()
			logf("engine %s: runtime install failed: %v", e.ID, err)
		}
		host.Unlock()
	}()
}

// EnsureRuntime makes sure the runtime is installed (installing it if need be) and waits for it; a job that arrives on a
// machine whose engines were not set up yet waits here instead of failing.
func (h hostRunner) EnsureRuntime(ctx context.Context, e modules.Engine, report func(string, float64)) error {
	if h.ready(e.ID) {
		return nil
	}
	h.kick(e)
	for {
		host.Lock()
		in := host.installs[e.ID]
		active, step, msg := in != nil && in.active, "", ""
		if in != nil {
			step, msg = in.step, in.err
		}
		host.Unlock()
		if h.ready(e.ID) {
			return nil
		}
		if !active {
			if msg == "" {
				msg = "the runtime is not installed"
			}
			return errors.New(msg)
		}
		report("installing the engine runtime: "+step, -1)
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(3 * time.Second):
		}
	}
}

func (h hostRunner) Start(engine string) error {
	e, ok := h.cat.Engine(engine)
	if !ok {
		return fmt.Errorf("unknown engine")
	}
	if !h.ready(engine) {
		h.kick(e)
		return fmt.Errorf("the engine runtime is still being installed on this machine (it needs to download Python and PyTorch, which takes a while; see the module page)")
	}
	p := h.paths(engine)
	argv, dir := engineArgv(h.cat, e, p)
	if len(argv) == 0 {
		return fmt.Errorf("no way to start %s directly on the machine", engine)
	}
	h.stopLegacyContainer(engine) // (can take a while: not under the lock)
	host.Lock()
	defer host.Unlock()
	if pr := host.procs[engine]; pr != nil {
		select {
		case <-pr.done:
		default:
			return nil // already running
		}
	}
	killStale(p.pid, p.venv)
	port := h.cat.HostPort(engine)
	if c, err := net.DialTimeout("tcp", "127.0.0.1:"+strconv.Itoa(port), 500*time.Millisecond); err == nil {
		c.Close()
		return fmt.Errorf("port %d is already used by another program", port)
	}
	for _, d := range []string{p.home, p.tmp} {
		if err := os.MkdirAll(d, 0755); err != nil {
			return err
		}
	}
	if fi, err := os.Stat(p.log); err == nil && fi.Size() > hostLogMax {
		os.Truncate(p.log, 0)
	}
	lf, err := os.OpenFile(p.log, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
	if err != nil {
		return err
	}
	fmt.Fprintf(lf, "\n=== %s starting %s ===\n", time.Now().Format(time.RFC3339), strings.Join(argv, " "))
	cmd := exec.Command(argv[0], argv[1:]...)
	cmd.Dir, cmd.Env, cmd.Stdout, cmd.Stderr = dir, engineEnv(p), lf, lf
	setProcGroup(cmd)
	if err := cmd.Start(); err != nil {
		lf.Close()
		return err
	}
	pr := &hostProc{cmd: cmd, pgid: cmd.Process.Pid, done: make(chan struct{}), started: time.Now()}
	host.procs[engine] = pr
	os.WriteFile(p.pid, []byte(strconv.Itoa(pr.pgid)+"\n"), 0644)
	go func() {
		err := cmd.Wait()
		fmt.Fprintf(lf, "=== %s exited: %v ===\n", time.Now().Format(time.RFC3339), err)
		lf.Close()
		os.Remove(p.pid)
		close(pr.done)
	}()
	return nil
}

// stopLegacyContainer removes the container an earlier version (or the Docker mode) left for this engine: it would hold
// the port. Without a runner on the machine there is nothing to remove.
func (h hostRunner) stopLegacyContainer(engine string) {
	sock := runnerSocket(h.cfg)
	if _, err := os.Stat(sock); err != nil {
		return
	}
	socketModRunner{runnerControl{sock: sock}}.Stop(engine)
}

func (h hostRunner) Stop(engine string) error {
	p := h.paths(engine)
	host.Lock()
	pr := host.procs[engine]
	host.Unlock()
	if pr != nil {
		stopGroup(pr)
	}
	killStale(p.pid, p.venv)
	return nil
}

// stopGroup asks the whole process group to end, and ends it by force when it has not within ten seconds.
func stopGroup(pr *hostProc) {
	select {
	case <-pr.done:
		return
	default:
	}
	signalGroup(pr.pgid, false)
	select {
	case <-pr.done:
	case <-time.After(10 * time.Second):
		signalGroup(pr.pgid, true)
		select {
		case <-pr.done:
		case <-time.After(5 * time.Second):
		}
	}
}

// killStale ends an engine left over by an agent that was killed without stopping it (its group id is in the pid file).
// It only acts on a process that really runs this engine's Python.
func killStale(pidFile, venv string) {
	b, err := os.ReadFile(pidFile)
	if err != nil {
		return
	}
	pid, _ := strconv.Atoi(strings.TrimSpace(string(b)))
	if pid > 1 && strings.Contains(procCmdline(pid), hostPython(venv)) {
		signalGroup(pid, true)
	}
	os.Remove(pidFile)
}

func (h hostRunner) Logs(engine string, tail int) string {
	p := h.paths(engine)
	f, err := os.Open(p.log)
	if err != nil {
		if in := func() *hostInstall { host.Lock(); defer host.Unlock(); return host.installs[engine] }(); in != nil && in.active {
			return "installing the engine runtime: " + in.step
		}
		return "(no log yet)"
	}
	defer f.Close()
	fi, _ := f.Stat()
	n := int64(64 << 10)
	if fi.Size() < n {
		n = fi.Size()
	}
	buf := make([]byte, n)
	f.ReadAt(buf, fi.Size()-n)
	lines := strings.Split(strings.TrimRight(string(buf), "\n"), "\n")
	if tail > 0 && len(lines) > tail {
		lines = lines[len(lines)-tail:]
	}
	return strings.Join(lines, "\n")
}

// Rmi removes the engine's Python environment (never anything outside the engine folder that carries our marker), and the
// shared uv folder once no engine has an environment any more.
func (h hostRunner) Rmi(engine string) error {
	p := h.paths(engine)
	if _, err := os.Stat(filepath.Join(p.dir, moduleMarker)); err != nil {
		return nil // not ours: nothing to remove
	}
	h.Stop(engine)
	if err := os.RemoveAll(p.venv); err != nil {
		return err
	}
	for _, d := range []string{p.home, p.tmp} {
		os.RemoveAll(d)
	}
	host.Lock()
	delete(host.installs, engine)
	host.Unlock()
	for _, e := range h.cat.Engines {
		if _, err := os.Stat(h.paths(e.ID).venv); err == nil {
			return nil
		}
	}
	return os.RemoveAll(p.hostDir)
}

// hostPauseAll stops (or continues) every engine we run directly: the load limits use it to slow them down.
func hostPauseAll(on bool) {
	host.Lock()
	defer host.Unlock()
	for _, pr := range host.procs {
		select {
		case <-pr.done:
		default:
			pauseGroup(pr.pgid, on)
		}
	}
}

// inOurHostEngine says whether a process belongs to one of the engines we run directly (for the Owner Guard: they are the
// machine serving other people, not its owner at work).
func inOurHostEngine(pid int) bool {
	g := pgidOf(pid)
	if g <= 0 {
		return false
	}
	host.Lock()
	defer host.Unlock()
	for _, pr := range host.procs {
		if pr.pgid == g {
			select {
			case <-pr.done:
			default:
				return true
			}
		}
	}
	return false
}

// ---- installing the runtime ----

// the commands and the uv download are variables so that tests can see the order without installing anything
var (
	runLoggedFn  = runLogged
	hostEnsureUV = ensureUV
)

func (h hostRunner) provision(ctx context.Context, e modules.Engine, step func(string)) error {
	p := h.paths(e.ID)
	cfg := h.cfg
	step("preparing the engine folder")
	if err := prepareEngine(cfg, e); err != nil {
		return err
	}
	for _, d := range []string{p.home, p.tmp, p.hostDir} {
		if err := os.MkdirAll(d, 0755); err != nil {
			return err
		}
	}
	lf, err := os.OpenFile(p.install, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
	if err != nil {
		return err
	}
	defer lf.Close()
	run := func(name string, argv ...string) (string, error) {
		return runLoggedFn(ctx, lf, p.dir, installEnv(p), name, argv...)
	}

	step("downloading uv " + uvVersion)
	uv, err := hostEnsureUV(ctx, p.hostDir)
	if err != nil {
		return err
	}
	py := hostPython(p.venv)
	if _, err := os.Stat(py); err != nil {
		step("installing Python 3.12")
		os.RemoveAll(p.venv)
		if _, err := run(uv, "venv", "--python", "3.12", p.venv); err != nil {
			return fmt.Errorf("python: %w", err)
		}
	}
	step("installing PyTorch (about 3 GB)")
	if _, err := run(uv, "pip", "install", "--python", py, "--torch-backend", hostTorchBk, "torch", "torchvision", "torchaudio"); err != nil {
		return fmt.Errorf("pytorch: %w", err)
	}
	// the engine's own packages must not replace the PyTorch that was just installed: hold the three at their versions
	frozen, err := run(uv, "pip", "freeze", "--python", py)
	if err != nil {
		return fmt.Errorf("pip freeze: %w", err)
	}
	var pins []string
	for _, l := range strings.Split(frozen, "\n") {
		if torchReqRe.MatchString(strings.TrimSpace(l)) {
			pins = append(pins, strings.TrimSpace(l))
		}
	}
	cons := filepath.Join(p.dir, "constraints.txt")
	if err := os.WriteFile(cons, []byte(strings.Join(pins, "\n")+"\n"), 0644); err != nil {
		return err
	}
	step("installing the engine's packages")
	switch e.ID {
	case "comfyui":
		req := filepath.Join(p.dir, filepath.Base(e.Workdir), "requirements.lab.txt")
		if _, err := run(uv, "pip", "install", "--python", py, "-r", req, "-c", cons); err != nil {
			return fmt.Errorf("packages: %w", err)
		}
	case "whisper":
		if _, err := run(uv, "pip", "install", "--python", py, "-c", cons, "transformers", "accelerate", "av", "numpy"); err != nil {
			return fmt.Errorf("packages: %w", err)
		}
	}
	step("checking the GPU")
	check := "import torch\nassert torch.cuda.is_available(), 'CUDA is not available to PyTorch'\n" +
		"a = torch.randn(256, 256, device='cuda'); (a @ a).sum().item(); torch.cuda.synchronize()\n" +
		"print('torch', torch.__version__, 'on', torch.cuda.get_device_name(0), torch.cuda.get_arch_list())\n"
	out, err := runLoggedFn(ctx, lf, p.dir, engineEnv(p), py, "-c", check)
	if err != nil {
		return fmt.Errorf("the GPU check failed (PyTorch cannot use the GPU of this machine; the Docker mode still works): %w", err)
	}
	return os.WriteFile(filepath.Join(p.venv, hostReady), []byte(strings.TrimSpace(out)+"\n"), 0644)
}

// runLogged runs a command, appends its output to the install log and returns it; an error carries the last lines.
func runLogged(ctx context.Context, lf io.Writer, dir string, env []string, name string, argv ...string) (string, error) {
	fmt.Fprintf(lf, "\n$ %s %s\n", name, strings.Join(argv, " "))
	var buf bytes.Buffer
	cmd := exec.CommandContext(ctx, name, argv...)
	cmd.Dir, cmd.Env = dir, env
	cmd.Stdout, cmd.Stderr = io.MultiWriter(&buf, lf), io.MultiWriter(&buf, lf)
	if err := cmd.Run(); err != nil {
		lines := strings.Split(strings.TrimSpace(buf.String()), "\n")
		if len(lines) > 6 {
			lines = lines[len(lines)-6:]
		}
		return buf.String(), fmt.Errorf("%v: %s", err, strings.Join(lines, " | "))
	}
	return buf.String(), nil
}

// ensureUV puts the pinned uv in hostDir/uv: downloaded once, checked against the pinned checksum before it is unpacked.
func ensureUV(ctx context.Context, hostDir string) (string, error) {
	dst := filepath.Join(hostDir, "uv")
	if out, err := exec.CommandContext(ctx, dst, "--version").Output(); err == nil && strings.Contains(string(out), uvVersion) {
		return dst, nil
	}
	sum, ok := uvSums[runtime.GOARCH]
	if !ok {
		return "", fmt.Errorf("no uv build is pinned for %s", runtime.GOARCH)
	}
	req, _ := http.NewRequestWithContext(ctx, "GET", uvURL(runtime.GOARCH), nil)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return "", fmt.Errorf("downloading uv: HTTP %d", resp.StatusCode)
	}
	data, err := io.ReadAll(io.LimitReader(resp.Body, 100<<20))
	if err != nil {
		return "", err
	}
	got := sha256.Sum256(data)
	if hex.EncodeToString(got[:]) != sum {
		return "", fmt.Errorf("the uv download does not match the pinned checksum; refusing to run it")
	}
	gz, err := gzip.NewReader(bytes.NewReader(data))
	if err != nil {
		return "", err
	}
	tr := tar.NewReader(gz)
	for {
		h, err := tr.Next()
		if err == io.EOF {
			return "", fmt.Errorf("the uv archive has no uv program")
		}
		if err != nil {
			return "", err
		}
		if h.Typeflag != tar.TypeReg || filepath.Base(h.Name) != "uv" {
			continue
		}
		tmp := dst + ".new"
		f, err := os.OpenFile(tmp, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0755)
		if err != nil {
			return "", err
		}
		if _, err := io.Copy(f, io.LimitReader(tr, 200<<20)); err != nil {
			f.Close()
			return "", err
		}
		f.Close()
		return dst, os.Rename(tmp, dst)
	}
}