package main

import (
	"archive/tar"
	"compress/gzip"
	"context"
	_ "embed"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
	"sync"
	"time"

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

// Optional AI modules (image, video, speech). A machine installs the ones its owner picks: the engine image (pulled by
// wxch-runner), the engine's code, and the model files. Nothing is installed by default. See internal/modules for the
// catalog; this file does the installing, moduleManagerjobs.go runs jobs on what is installed.

//go:embed whisper_server.py
var whisperServerPy []byte

const moduleMarker = ".wxch-module"

// loadCatalog is a variable so tests can use a small catalog that points at a fake download server.
var loadCatalog = modules.Load

// modRunner is what the manager needs from wxch-runner; tests replace it.
type modRunner interface {
	Status() ([]runnerModule, error)
	Pull(engine string) error
	Start(engine string) error
	Stop(engine string) error
	Logs(engine string, tail int) string
	Rmi(engine string) error
}

type runnerModule struct {
	Engine       string `json:"engine"`
	State        string `json:"state"`
	ImagePresent bool   `json:"image_present"`
	Pulling      bool   `json:"pulling"`
	PullErr      string `json:"pull_err"`
	PullStep     string `json:"pull_step,omitempty"` // what the runtime install is doing now (engines run directly on the machine)
	Port         int    `json:"port"`
	Pinned       bool   `json:"pinned"` // started by hand: not stopped when idle (set by the agent, not the runner)
}

type socketModRunner struct{ rc runnerControl }

func (s socketModRunner) Status() ([]runnerModule, error) {
	d, err := s.rc.call(runnerRequest{Op: "mod_status"}, 60*time.Second)
	if err != nil {
		return nil, err
	}
	var out []runnerModule
	return out, json.Unmarshal(d, &out)
}
func (s socketModRunner) Pull(e string) error {
	_, err := s.rc.call(runnerRequest{Op: "mod_pull", Engine: e}, 20*time.Second)
	return err
}
func (s socketModRunner) Start(e string) error {
	_, err := s.rc.call(runnerRequest{Op: "mod_start", Engine: e}, 3*time.Minute)
	return err
}
func (s socketModRunner) Stop(e string) error {
	_, err := s.rc.call(runnerRequest{Op: "mod_stop", Engine: e}, 90*time.Second)
	return err
}
func (s socketModRunner) Logs(e string, tail int) string {
	d, err := s.rc.call(runnerRequest{Op: "mod_logs", Engine: e, Tail: tail}, 20*time.Second)
	if err != nil {
		return err.Error()
	}
	var out string
	json.Unmarshal(d, &out)
	return out
}
func (s socketModRunner) Rmi(e string) error {
	_, err := s.rc.call(runnerRequest{Op: "mod_rmi", Engine: e}, 6*time.Minute)
	return err
}

var modRunnerOverride modRunner

// engineMode says how the engines run on this machine: "host" (a Python environment of their own, started by the agent;
// the default) or "docker" (containers started by wxch-runner). state.json's mode wins (the agent can write it), then
// engine_mode in the node config, which the installer sets.
func engineMode(cfg nodeConfig) string {
	for _, m := range []string{readModuleState(cfg).Mode, cfg.EngineMode} {
		if m == "host" || m == "docker" {
			return m
		}
	}
	return "host"
}

// runtimeInstaller is implemented by the runner of engines that run directly on the machine: it installs the Python
// environment on demand, so a job that arrives before it is there waits instead of failing.
type runtimeInstaller interface {
	EnsureRuntime(ctx context.Context, e modules.Engine, report func(string, float64)) error
}

func pickModRunner(cfg nodeConfig) (modRunner, error) {
	if modRunnerOverride != nil {
		return modRunnerOverride, nil
	}
	if engineMode(cfg) == "host" {
		cat, err := loadCatalog()
		if err != nil {
			return nil, err
		}
		return hostRunner{cfg: cfg, cat: cat}, nil
	}
	if _, err := os.Stat(runnerSocket(cfg)); err != nil {
		return nil, fmt.Errorf("the runner is not installed on this machine: %s", installHint())
	}
	return socketModRunner{runnerControl{sock: runnerSocket(cfg)}}, nil
}

// ---- state ----

type moduleState struct {
	Installed map[string]string `json:"installed"`        // capability id -> when
	Pinned    map[string]bool   `json:"pinned,omitempty"` // engines the administrator started by hand: kept until stopped by hand
	Mode      string            `json:"mode,omitempty"`   // "host" or "docker": how the engines run on this machine (see engineMode)
}

// setPinned remembers (or forgets) that an engine was started by hand and must not be stopped as idle.
func setPinned(cfg nodeConfig, engine string, on bool) {
	err := updateModuleState(cfg, func(st *moduleState) {
		if on {
			if st.Pinned == nil {
				st.Pinned = map[string]bool{}
			}
			st.Pinned[engine] = true
		} else {
			delete(st.Pinned, engine)
		}
	})
	if err != nil {
		logf("module state: %v", err)
	}
}

type installProgress struct {
	ID      string `json:"id"`
	Op      string `json:"op"` // install or remove
	Status  string `json:"status"`
	Total   int64  `json:"total"`
	Done    int64  `json:"done"`
	Err     string `json:"err,omitempty"`
	Started int64  `json:"started"`
	Ended   int64  `json:"ended,omitempty"`
}

type moduleMgr struct {
	mu   sync.Mutex
	busy *installProgress // the one install or removal in progress
	last map[string]*installProgress
}

var moduleManager = &moduleMgr{last: map[string]*installProgress{}}

func modulesDir(cfg nodeConfig) string {
	if cfg.ModulesDir != "" {
		return cfg.ModulesDir
	}
	return "/mnt/models/modules"
}

// stateMu serialises every change of state.json: pinning, finishing an install, finishing a removal and the mode all read,
// change and write the same file, and without the lock one of them could undo another.
var stateMu sync.Mutex

func statePath(cfg nodeConfig) string { return filepath.Join(modulesDir(cfg), "state.json") }

// loadState reads state.json. A missing file is an empty state; a file that cannot be decoded is an error, so that it is
// never overwritten by a guess (the last good copy is kept next to it as state.json.bak).
func loadState(cfg nodeConfig) (moduleState, error) {
	st := moduleState{Installed: map[string]string{}}
	b, err := os.ReadFile(statePath(cfg))
	if err != nil {
		if os.IsNotExist(err) {
			return st, nil
		}
		return st, err
	}
	if err := json.Unmarshal(b, &st); err != nil {
		return moduleState{Installed: map[string]string{}}, fmt.Errorf("%s cannot be read (%v)", statePath(cfg), err)
	}
	if st.Installed == nil {
		st.Installed = map[string]string{}
	}
	return st, nil
}

var stateWarned sync.Once

// readModuleState is for looking: a damaged file shows as empty (and is logged once); changes go through updateModuleState,
// which refuses to touch a damaged file.
func readModuleState(cfg nodeConfig) moduleState {
	st, err := loadState(cfg)
	if err != nil {
		stateWarned.Do(func() { logf("module state: %v", err) })
	}
	return st
}

// saveState writes atomically through a temp file of its own name, keeping the previous good file as state.json.bak.
func saveState(cfg nodeConfig, st moduleState) error {
	b, _ := json.MarshalIndent(st, "", " ")
	dir := modulesDir(cfg)
	tmp, err := os.CreateTemp(dir, "state-*.tmp")
	if err != nil {
		return err
	}
	if _, err := tmp.Write(b); err != nil {
		tmp.Close()
		os.Remove(tmp.Name())
		return err
	}
	if err := tmp.Close(); err != nil {
		os.Remove(tmp.Name())
		return err
	}
	os.Chmod(tmp.Name(), 0644)
	if old, err := os.ReadFile(statePath(cfg)); err == nil && json.Valid(old) {
		os.WriteFile(statePath(cfg)+".bak", old, 0644)
	}
	if err := os.Rename(tmp.Name(), statePath(cfg)); err != nil {
		os.Remove(tmp.Name())
		return err
	}
	return nil
}

// updateModuleState is the only way to change the state: lock, read, change, write.
func updateModuleState(cfg nodeConfig, fn func(*moduleState)) error {
	stateMu.Lock()
	defer stateMu.Unlock()
	st, err := loadState(cfg)
	if err != nil {
		return fmt.Errorf("the state file is damaged, fix or remove it (the last good copy is state.json.bak): %v", err)
	}
	fn(&st)
	return saveState(cfg, st)
}

// writeModuleState replaces the whole state (tests and first set-up).
func writeModuleState(cfg nodeConfig, st moduleState) error {
	stateMu.Lock()
	defer stateMu.Unlock()
	return saveState(cfg, st)
}

func engineDir(cfg nodeConfig, e modules.Engine) string { return filepath.Join(modulesDir(cfg), e.ID) }

func (m *moduleMgr) setBusy(p *installProgress) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.busy != nil {
		return fmt.Errorf("another module task is running (%s %s); wait for it to finish", m.busy.Op, m.busy.ID)
	}
	m.busy = p
	return nil
}

func (m *moduleMgr) finish(p *installProgress, err error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	p.Ended = time.Now().Unix()
	if err != nil {
		p.Err = err.Error()
		p.Status = "failed"
	} else {
		p.Status = "done"
	}
	m.last[p.ID] = p
	m.busy = nil
}

func (m *moduleMgr) update(p *installProgress, status string, total, done int64) {
	m.mu.Lock()
	p.Status = status
	if total >= 0 {
		p.Total = total
	}
	if done >= 0 {
		p.Done = done
	}
	m.mu.Unlock()
}

// ---- install ----

var torchReqRe = regexp.MustCompile(`(?i)^(torch|torchvision|torchaudio)([=<>~! ]|$)`)
var gitTagRe = regexp.MustCompile(`^v?[0-9][0-9A-Za-z._-]{0,30}$`)

// capabilitiesUsing lists the installed capabilities that run on this engine.
func capabilitiesUsing(cat *modules.Catalog, st moduleState, engine string) []string {
	var out []string
	for _, c := range cat.Capabilities {
		if _, ok := st.Installed[c.ID]; ok && c.Engine == engine {
			out = append(out, c.ID)
		}
	}
	return out
}

func missingBytes(cfg nodeConfig, e modules.Engine, cp modules.Capability) int64 {
	var n int64
	for _, mf := range cp.Models {
		p := filepath.Join(engineDir(cfg, e), filepath.FromSlash(e.ModelsRoot), filepath.FromSlash(mf.Dest))
		if fi, err := os.Stat(p); err != nil || fi.Size() != mf.Bytes {
			n += mf.Bytes
		}
	}
	return n
}

// Install starts installing a capability in the background; progress shows in modules.list.
func (m *moduleMgr) Install(cfg nodeConfig, id string) error {
	cat, err := loadCatalog()
	if err != nil {
		return err
	}
	cp, ok := cat.Capability(id)
	if !ok {
		return fmt.Errorf("unknown capability %q", id)
	}
	e, _ := cat.Engine(cp.Engine)
	if _, done := readModuleState(cfg).Installed[id]; done {
		return fmt.Errorf("%s is already installed", id)
	}
	rn, err := pickModRunner(cfg)
	if err != nil {
		return err
	}
	if err := os.MkdirAll(modulesDir(cfg), 0755); err != nil {
		return fmt.Errorf("cannot create %s (run the installer's models step): %v", modulesDir(cfg), err)
	}
	// the space the install still needs: files not yet there, the engine folder, and the image if it is not pulled
	need := missingBytes(cfg, e, cp)
	if _, err := os.Stat(filepath.Join(engineDir(cfg, e), moduleMarker)); err != nil {
		need += e.WorkBytes
	}
	if st, err := rn.Status(); err == nil {
		have := false
		for _, s := range st {
			if s.Engine == e.ID && s.ImagePresent {
				have = true
			}
		}
		if !have {
			if engineMode(cfg) == "host" {
				need += e.HostBytes
			} else {
				need += e.ImageBytes
			}
		}
	}
	if free, _, err := diskFree(modulesDir(cfg)); err == nil && uint64(need)+uint64(cfg.MinFreeGB)<<30 > free {
		return fmt.Errorf("not enough free disk space: this needs about %.1f GB (+%d GB spare), %.1f GB is free", float64(need)/1e9, cfg.MinFreeGB, float64(free)/1e9)
	}
	p := &installProgress{ID: id, Op: "install", Status: "starting", Total: need, Started: time.Now().Unix()}
	if err := m.setBusy(p); err != nil {
		return err
	}
	go func() {
		m.finish(p, m.doInstall(context.Background(), cfg, cat, cp, e, rn, p))
	}()
	return nil
}

func (m *moduleMgr) doInstall(ctx context.Context, cfg nodeConfig, cat *modules.Catalog, cp modules.Capability, e modules.Engine, rn modRunner, p *installProgress) error {
	// 1. the engine image (the runner downloads it; it reports no byte counts, only whether it is there), or, for engines
	// that run directly on the machine, their Python environment
	direct := engineMode(cfg) == "host"
	what := "downloading the engine image (large, can take a long time)"
	if direct {
		what = "installing the engine runtime (Python and PyTorch, about 5 GB)"
	}
	m.update(p, what, -1, -1)
	started := false // an error left by an earlier attempt is not this attempt's
	for {
		st, err := rn.Status()
		if err != nil {
			return err
		}
		var mine *runnerModule
		for i := range st {
			if st[i].Engine == e.ID {
				mine = &st[i]
			}
		}
		if mine == nil {
			return fmt.Errorf("the runner does not know the engine %s (update the runner)", e.ID)
		}
		if mine.ImagePresent {
			break
		}
		if started && mine.PullErr != "" && !mine.Pulling {
			return fmt.Errorf("the engine could not be set up: %s", mine.PullErr)
		}
		if !mine.Pulling {
			if err := rn.Pull(e.ID); err != nil {
				return err
			}
			started = true
			continue // look again at once: an image that is already there needs no waiting
		}
		if mine.PullStep != "" {
			m.update(p, what+": "+mine.PullStep, -1, -1)
		}
		time.Sleep(5 * time.Second)
	}
	// 2. the engine folder and its code
	m.update(p, "preparing the engine", -1, -1)
	if err := prepareEngine(cfg, e); err != nil {
		return err
	}
	// 3. the model files
	var done int64
	for _, mf := range cp.Models {
		dest := filepath.Join(engineDir(cfg, e), filepath.FromSlash(e.ModelsRoot), filepath.FromSlash(mf.Dest))
		if fi, err := os.Stat(dest); err == nil && fi.Size() == mf.Bytes {
			if mf.SHA256 == "" {
				continue
			}
			m.update(p, "checking "+mf.Dest, -1, -1)
		} else {
			m.update(p, "downloading "+mf.Dest, -1, -1)
		}
		before := done
		err := downloadFile(ctx, mf.URL, dest, mf.Bytes, mf.SHA256, func(n int64) {
			done += n
			m.update(p, "downloading "+mf.Dest, -1, done)
		})
		if err != nil {
			return err
		}
		done = before + mf.Bytes
		m.update(p, "downloading "+mf.Dest, -1, done)
	}
	return updateModuleState(cfg, func(st *moduleState) {
		st.Installed[cp.ID] = time.Now().UTC().Format(time.RFC3339)
	})
}

var gitLookPath = exec.LookPath

// codeArchiveURL is where the tagged source of an engine can be downloaded as an archive (a variable for tests).
var codeArchiveURL = func(e modules.Engine) string {
	return strings.TrimSuffix(e.SourceRepo, "/") + "/archive/refs/tags/" + e.SourceTag + ".tar.gz"
}

// fetchCodeArchive downloads the tagged source as a tar.gz and unpacks it into dest, dropping the top folder of the
// archive. Only plain files and folders are created, never outside dest.
func fetchCodeArchive(e modules.Engine, dest string) error {
	resp, err := (&http.Client{Timeout: 10 * time.Minute}).Get(codeArchiveURL(e))
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return fmt.Errorf("HTTP %d", resp.StatusCode)
	}
	gz, err := gzip.NewReader(io.LimitReader(resp.Body, 300<<20))
	if err != nil {
		return err
	}
	defer gz.Close()
	tmp := dest + ".unpack"
	os.RemoveAll(tmp)
	if err := os.MkdirAll(tmp, 0755); err != nil {
		return err
	}
	defer os.RemoveAll(tmp)
	tr := tar.NewReader(gz)
	var total int64
	for {
		h, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		_, rel, ok := strings.Cut(strings.TrimPrefix(filepath.ToSlash(h.Name), "./"), "/")
		if !ok || rel == "" || strings.Contains(rel, "..") || strings.HasPrefix(rel, "/") {
			continue // the top folder itself, a global header, or something that should not be there
		}
		target := filepath.Join(tmp, filepath.FromSlash(rel))
		switch h.Typeflag {
		case tar.TypeDir:
			os.MkdirAll(target, 0755)
		case tar.TypeReg:
			if total += h.Size; total > 400<<20 {
				return fmt.Errorf("the archive is larger than expected")
			}
			os.MkdirAll(filepath.Dir(target), 0755)
			mode := os.FileMode(0644)
			if h.Mode&0111 != 0 {
				mode = 0755
			}
			f, err := os.OpenFile(target, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, mode)
			if err != nil {
				return err
			}
			if _, err := io.Copy(f, tr); err != nil {
				f.Close()
				return err
			}
			f.Close()
		}
	}
	if _, err := os.Stat(filepath.Join(tmp, "main.py")); err != nil {
		return fmt.Errorf("the archive does not hold the engine")
	}
	return os.Rename(tmp, dest)
}

// prepareEngine creates the engine folder (marked as ours) and puts the engine's code in it.
func prepareEngine(cfg nodeConfig, e modules.Engine) error {
	dir := engineDir(cfg, e)
	if err := os.MkdirAll(dir, 0755); err != nil {
		return err
	}
	if err := os.WriteFile(filepath.Join(dir, moduleMarker), []byte(e.ID+"\n"), 0644); err != nil {
		return err
	}
	if e.SourceRepo != "" {
		src := filepath.Join(dir, filepath.Base(e.Workdir))
		if _, err := os.Stat(filepath.Join(src, "main.py")); err != nil {
			if !gitTagRe.MatchString(e.SourceTag) || !strings.HasPrefix(e.SourceRepo, "https://github.com/") {
				return fmt.Errorf("bad engine source in the catalog")
			}
			os.RemoveAll(src) // a half-finished clone
			if _, lerr := gitLookPath("git"); lerr == nil {
				out, err := exec.Command("git", "clone", "-q", "--depth", "1", "--branch", e.SourceTag, e.SourceRepo, src).CombinedOutput()
				if err != nil {
					os.RemoveAll(src)
					return fmt.Errorf("could not fetch the engine code: %v: %s", err, strings.TrimSpace(string(out)))
				}
			} else if err := fetchCodeArchive(e, src); err != nil { // no git on this machine: the tagged archive does as well
				os.RemoveAll(src)
				return fmt.Errorf("could not fetch the engine code (git is not installed, the archive failed too): %v", err)
			}
		}
		// the container brings its own torch (built for this machine's GPU); installing the pinned one from pip would replace it
		req, err := os.ReadFile(filepath.Join(src, "requirements.txt"))
		if err != nil {
			return err
		}
		var keep []string
		for _, l := range strings.Split(string(req), "\n") {
			if torchReqRe.MatchString(strings.TrimSpace(l)) {
				continue
			}
			keep = append(keep, strings.TrimRight(l, "\r"))
		}
		if err := os.WriteFile(filepath.Join(src, "requirements.lab.txt"), []byte(strings.Join(keep, "\n")+"\n"), 0644); err != nil {
			return err
		}
		return os.MkdirAll(filepath.Join(src, "models"), 0755)
	}
	if e.ID == "whisper" {
		return os.WriteFile(filepath.Join(dir, "whisper_server.py"), whisperServerPy, 0644)
	}
	return nil
}

// ---- remove ----

// Remove deletes a capability's model files that no other installed capability needs. When it was the last capability of
// its engine, the engine folder goes too (only a folder carrying our marker), and the image when no engine needs it.
func (m *moduleMgr) Remove(cfg nodeConfig, id string) error {
	cat, err := loadCatalog()
	if err != nil {
		return err
	}
	cp, ok := cat.Capability(id)
	if !ok {
		return fmt.Errorf("unknown capability %q", id)
	}
	st := readModuleState(cfg)
	if _, ok := st.Installed[id]; !ok {
		return fmt.Errorf("%s is not installed", id)
	}
	if jobsBusy(id) {
		return fmt.Errorf("a job is running; wait for it to finish")
	}
	rn, err := pickModRunner(cfg)
	if err != nil {
		return err
	}
	p := &installProgress{ID: id, Op: "remove", Status: "removing", Started: time.Now().Unix()}
	if err := m.setBusy(p); err != nil {
		return err
	}
	err = m.doRemove(cfg, cat, cp, st, rn)
	m.finish(p, err)
	return err
}

func (m *moduleMgr) doRemove(cfg nodeConfig, cat *modules.Catalog, cp modules.Capability, st moduleState, rn modRunner) error {
	e, _ := cat.Engine(cp.Engine)
	// what will be left once this capability is gone; the state file itself is changed only after the cleanup worked
	remaining := moduleState{Installed: map[string]string{}, Pinned: st.Pinned}
	for k, v := range st.Installed {
		if k != cp.ID {
			remaining.Installed[k] = v
		}
	}
	commit := func() error {
		return updateModuleState(cfg, func(cur *moduleState) { delete(cur.Installed, cp.ID) })
	}
	stillNeeded := map[string]bool{}
	for _, other := range cat.Capabilities {
		if _, ok := remaining.Installed[other.ID]; ok && other.Engine == e.ID {
			for _, mf := range other.Models {
				stillNeeded[mf.Dest] = true
			}
		}
	}
	root := filepath.Join(engineDir(cfg, e), filepath.FromSlash(e.ModelsRoot))
	for _, mf := range cp.Models {
		if stillNeeded[mf.Dest] {
			continue
		}
		p := filepath.Join(root, filepath.FromSlash(mf.Dest))
		if !strings.HasPrefix(p, root+string(os.PathSeparator)) {
			continue
		}
		if err := os.Remove(p); err != nil && !os.IsNotExist(err) {
			return fmt.Errorf("could not delete %s: %v", mf.Dest, err)
		}
		os.Remove(p + ".part")
		for d := filepath.Dir(p); d != root && strings.HasPrefix(d, root); d = filepath.Dir(d) {
			if os.Remove(d) != nil { // only an empty folder goes
				break
			}
		}
	}
	if len(capabilitiesUsing(cat, remaining, e.ID)) > 0 {
		return commit()
	}
	rn.Stop(e.ID)
	dir := engineDir(cfg, e)
	if _, err := os.Stat(filepath.Join(dir, moduleMarker)); err == nil {
		// the engine folder is ours, but its models folder may hold files the administrator added by hand (ComfyUI has
		// its own page for that): those are never deleted, and then the folder stays
		if extra := foreignModelFiles(cfg, cat, e); len(extra) > 0 {
			logf("module remove: %s keeps its folder, it holds files that are not ours (%s)", e.ID, strings.Join(extra, ", "))
		} else if err := os.RemoveAll(dir); err != nil {
			return err
		}
	}
	if err := commit(); err != nil { // the files are gone: the state must say so
		return err
	}
	// the image is shared by engines with the same digest
	for _, other := range cat.Engines {
		if other.ID != e.ID && other.Digest == e.Digest && len(capabilitiesUsing(cat, remaining, other.ID)) > 0 {
			return nil
		}
	}
	return rn.Rmi(e.ID)
}

// extraFile is a file in an engine's models folder that no catalog entry accounts for: something the administrator put
// there (ComfyUI's own page can download models). The system did not install it and does not count it in the sizes of the
// modules, but the page lists it and can delete it.
type extraFile struct {
	Engine string `json:"engine"`
	Path   string `json:"path"` // relative to the engine's models folder, with /
	Size   int64  `json:"size"`
}

const maxExtraFiles = 300

// walkExtra finds the extra files of one engine (up to limit). ComfyUI ships placeholder files (put_*) and config
// samples (configs/) in its models folder, and a download in progress is a .part file: none of these count.
func walkExtra(cfg nodeConfig, cat *modules.Catalog, e modules.Engine, limit int) []extraFile {
	known := map[string]bool{}
	for _, cp := range cat.Capabilities {
		if cp.Engine == e.ID {
			for _, mf := range cp.Models {
				known[filepath.ToSlash(mf.Dest)] = true
			}
		}
	}
	root := filepath.Join(engineDir(cfg, e), filepath.FromSlash(e.ModelsRoot))
	var extra []extraFile
	filepath.WalkDir(root, func(p string, d os.DirEntry, err error) error {
		if err != nil || len(extra) >= limit {
			return nil
		}
		rel, rerr := filepath.Rel(root, p)
		if rerr != nil {
			return nil
		}
		rel = filepath.ToSlash(rel)
		if d.IsDir() {
			if rel == "configs" {
				return filepath.SkipDir
			}
			return nil
		}
		if known[rel] || strings.HasPrefix(d.Name(), "put_") || strings.HasSuffix(d.Name(), ".part") {
			return nil
		}
		var size int64
		if fi, err := os.Stat(p); err == nil { // follows a link to a file on an external disk
			size = fi.Size()
		}
		extra = append(extra, extraFile{Engine: e.ID, Path: rel, Size: size})
		return nil
	})
	return extra
}

// foreignModelFiles lists (a few of) the extra files of an engine; the folder is only deleted when there are none.
func foreignModelFiles(cfg nodeConfig, cat *modules.Catalog, e modules.Engine) []string {
	var out []string
	for _, f := range walkExtra(cfg, cat, e, 5) {
		out = append(out, f.Path)
	}
	return out
}

// listExtraFiles lists the extra files of every engine that has a folder.
func listExtraFiles(cfg nodeConfig, cat *modules.Catalog) []extraFile {
	out := []extraFile{}
	for _, e := range cat.Engines {
		if _, err := os.Stat(engineDir(cfg, e)); err != nil {
			continue
		}
		out = append(out, walkExtra(cfg, cat, e, maxExtraFiles)...)
	}
	return out
}

// DeleteExtra deletes one extra file the administrator chose. Only a file the listing shows can go (so never a file of a
// module, a placeholder or anything outside the models folder), and not while a job or an install is running.
func (m *moduleMgr) DeleteExtra(cfg nodeConfig, engine, rel string) error {
	cat, err := loadCatalog()
	if err != nil {
		return err
	}
	e, ok := cat.Engine(engine)
	if !ok {
		return fmt.Errorf("unknown engine")
	}
	rel = strings.TrimPrefix(filepath.ToSlash(rel), "./")
	if rel == "" || strings.HasPrefix(rel, "/") || strings.Contains(rel, "..") {
		return fmt.Errorf("invalid path")
	}
	found := false
	for _, f := range walkExtra(cfg, cat, e, maxExtraFiles*4) {
		if f.Path == rel {
			found = true
			break
		}
	}
	if !found {
		return fmt.Errorf("%s is not an added file (module files are removed by uninstalling the module)", rel)
	}
	if jobsBusy("") || tunnelsOpen() {
		return fmt.Errorf("a job is running or the ComfyUI page is open; try again when it is idle")
	}
	m.mu.Lock()
	busy := m.busy != nil
	m.mu.Unlock()
	if busy {
		return fmt.Errorf("a module task is running; try again when it has finished")
	}
	root := filepath.Join(engineDir(cfg, e), filepath.FromSlash(e.ModelsRoot))
	full := filepath.Join(root, filepath.FromSlash(rel))
	if !strings.HasPrefix(full, root+string(os.PathSeparator)) {
		return fmt.Errorf("invalid path")
	}
	if err := os.Remove(full); err != nil { // a link to another disk is removed as the link
		return err
	}
	for d := filepath.Dir(full); d != root && strings.HasPrefix(d, root+string(os.PathSeparator)); d = filepath.Dir(d) {
		if os.Remove(d) != nil { // only an empty folder goes
			break
		}
	}
	logf("module files: deleted %s/%s", engine, rel)
	return nil
}

// ---- list ----

type capStatus struct {
	ID          string           `json:"id"`
	Installed   bool             `json:"installed"`
	InstalledAt string           `json:"installed_at,omitempty"`
	Task        *installProgress `json:"task,omitempty"`
}

type modulesOverview struct {
	Dir          string          `json:"dir"`
	FreeBytes    uint64          `json:"free_bytes"`
	TotalBytes   uint64          `json:"total_bytes"`
	Capabilities []capStatus     `json:"capabilities"`
	Engines      []runnerModule  `json:"engines"`
	Mode         string          `json:"mode"` // how the engines run on this machine: host or docker
	RunnerError  string          `json:"runner_error,omitempty"`
	Updated      string          `json:"catalog_updated"`
	Jobs         []jobStatusInfo `json:"jobs,omitempty"`
	Extra        []extraFile     `json:"extra"` // files added by hand to an engine's models folder
}

func (m *moduleMgr) Overview(cfg nodeConfig) modulesOverview {
	cat, err := loadCatalog()
	if err != nil {
		return modulesOverview{RunnerError: err.Error()}
	}
	st := readModuleState(cfg)
	o := modulesOverview{Dir: modulesDir(cfg), Updated: cat.Updated, Mode: engineMode(cfg)}
	if free, total, err := diskFree(modulesDir(cfg)); err == nil {
		o.FreeBytes, o.TotalBytes = free, total
	}
	m.mu.Lock()
	for _, c := range cat.Capabilities {
		cs := capStatus{ID: c.ID}
		if at, ok := st.Installed[c.ID]; ok {
			cs.Installed, cs.InstalledAt = true, at
		}
		if m.busy != nil && m.busy.ID == c.ID {
			cp := *m.busy
			cs.Task = &cp
		} else if l := m.last[c.ID]; l != nil {
			cp := *l
			cs.Task = &cp
		}
		o.Capabilities = append(o.Capabilities, cs)
	}
	m.mu.Unlock()
	sort.Slice(o.Capabilities, func(i, j int) bool { return o.Capabilities[i].ID < o.Capabilities[j].ID })
	if rn, err := pickModRunner(cfg); err != nil {
		o.RunnerError = err.Error()
	} else if es, err := rn.Status(); err != nil {
		o.RunnerError = err.Error()
	} else {
		o.Engines = es
		for i := range o.Engines {
			o.Engines[i].Pinned = st.Pinned[o.Engines[i].Engine]
		}
	}
	o.Jobs = jobsOverview()
	o.Extra = listExtraFiles(cfg, cat)
	return o
}