package main

import (
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

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

// Module engines (ComfyUI, Whisper) run like vLLM: the agent asks, the runner decides. Everything that matters is fixed
// here from the compiled-in catalog: the image (pinned by digest), the port, the folder that is mounted and the command.
// The agent can only say which engine to start or stop, never what to run.

const moduleLabel = "wxch.module"

func moduleContainer(engine string) string { return "wxch-mod-" + engine }

type moduleInfo struct {
	Engine       string `json:"engine"`
	State        string `json:"state"` // running, exited, created, "" (no container)
	ImagePresent bool   `json:"image_present"`
	Pulling      bool   `json:"pulling"`
	PullErr      string `json:"pull_err,omitempty"`
	Port         int    `json:"port"`
}

var pulls = struct {
	sync.Mutex
	active map[string]bool
	errs   map[string]string
}{active: map[string]bool{}, errs: map[string]string{}}

// moduleDir returns the engine folder under the modules folder (or an extra root), resolved, and proves it is a folder.
func moduleDir(cfg Config, engine string) (string, error) {
	root, err := filepath.EvalSymlinks(cfg.ModulesDir)
	if err != nil {
		return "", fmt.Errorf("modules folder missing: %v", err)
	}
	dir, err := filepath.EvalSymlinks(filepath.Join(root, engine))
	if err != nil {
		return "", fmt.Errorf("the %s folder does not exist yet; install a module that uses it first", engine)
	}
	if st, err := os.Stat(dir); err != nil || !st.IsDir() {
		return "", fmt.Errorf("the %s folder is not a folder", engine)
	}
	if !strings.HasPrefix(dir, root+string(os.PathSeparator)) {
		allowed := false
		for _, e := range cfg.ExtraRoots {
			if er, err := filepath.EvalSymlinks(e); err == nil && strings.HasPrefix(dir, er+string(os.PathSeparator)) {
				allowed = true
			}
		}
		if !allowed {
			return "", fmt.Errorf("the %s folder points outside the allowed folders", engine)
		}
	}
	return dir, nil
}

// buildModuleRun returns the exact docker arguments for an engine.
func buildModuleRun(cfg Config, cat *modules.Catalog, engine string) ([]string, error) {
	e, ok := cat.Engine(engine)
	if !ok {
		return nil, fmt.Errorf("unknown engine")
	}
	dir, err := moduleDir(cfg, engine)
	if err != nil {
		return nil, err
	}
	gid := cfg.AgentGID
	if gid < 0 {
		gid = cfg.AgentUID
	}
	args := []string{"run", "-d", "--name", moduleContainer(engine),
		"--label", moduleLabel + "=" + engine,
		"--gpus", "all", "--shm-size", cfg.ShmSize, "--ulimit", "memlock=-1", "--ulimit", "stack=67108864",
		// ComfyUI with PyTorch and CUDA keeps many files open (models, sockets, the GPU devices); with Docker's low default
		// (a soft limit of 1024 on some installs) jobs fail with "Too many open files" or odd CUDA errors
		"--ulimit", "nofile=1048576:1048576",
		"--security-opt", "no-new-privileges", "--cap-drop", "ALL",
		"--user", fmt.Sprintf("%d:%d", cfg.AgentUID, gid),
		"-e", "HOME=/work", "-e", "PYTHONUSERBASE=/work/pyuser", "-e", "DO_NOT_TRACK=1", "-e", "HF_HUB_OFFLINE=1",
		"-p", fmt.Sprintf("127.0.0.1:%d:%d", cat.HostPort(engine), e.Port),
		"-v", dir + ":/work", "-w", e.Workdir}
	// models moved to the external disk are links out of the engine folder: the folders root listed are visible at the
	// same path, read-only, so the links resolve inside the container
	for _, x := range cfg.ExtraRoots {
		if er, err := filepath.EvalSymlinks(x); err == nil {
			args = append(args, "-v", er+":"+er+":ro")
		}
	}
	args = append(args, e.ImageRef())
	return append(args, e.Cmd...), nil
}

func moduleStatus(cfg Config, cat *modules.Catalog) []moduleInfo {
	states := map[string]string{}
	if out, err := docker(10*time.Second, "ps", "-a", "--filter", "label="+moduleLabel, "--format", `{{.Label "`+moduleLabel+`"}}|{{.State}}`); err == nil {
		for _, line := range strings.Split(strings.TrimSpace(out), "\n") {
			if a, b, ok := strings.Cut(line, "|"); ok {
				states[a] = b
			}
		}
	}
	pulls.Lock()
	defer pulls.Unlock()
	var out []moduleInfo
	for _, e := range cat.Engines {
		_, err := docker(8*time.Second, "image", "inspect", "--format", "{{.Id}}", e.ImageRef())
		out = append(out, moduleInfo{Engine: e.ID, State: states[e.ID], ImagePresent: err == nil, Pulling: pulls.active[e.Image+e.Digest], PullErr: pulls.errs[e.Image+e.Digest], Port: cat.HostPort(e.ID)})
	}
	return out
}

// startPull downloads the engine image in the background (it is tens of gigabytes); moduleStatus tells how it went.
func startPull(e modules.Engine) error {
	key := e.Image + e.Digest
	pulls.Lock()
	if pulls.active[key] {
		pulls.Unlock()
		return nil
	}
	pulls.active[key] = true
	delete(pulls.errs, key)
	pulls.Unlock()
	go func() {
		_, err := docker(3*time.Hour, "pull", e.ImageRef())
		pulls.Lock()
		delete(pulls.active, key)
		if err != nil {
			pulls.errs[key] = err.Error()
		}
		pulls.Unlock()
	}()
	return nil
}

var moduleCatalog = func() *modules.Catalog {
	c, err := modules.Load()
	if err != nil {
		panic(err) // a broken catalog is a build mistake, caught by the tests
	}
	return c
}()

// handleModule serves the mod_* operations.
func handleModule(cfg Config, req Request) Response {
	cat := moduleCatalog
	if req.Op == "mod_status" {
		return ok(moduleStatus(cfg, cat))
	}
	e, found := cat.Engine(req.Engine)
	if !found {
		return fail("unknown engine")
	}
	switch req.Op {
	case "mod_pull":
		startPull(e)
		return ok(nil)
	case "mod_start":
		if _, err := docker(8*time.Second, "image", "inspect", "--format", "{{.Id}}", e.ImageRef()); err != nil {
			return fail("the engine image is not downloaded yet")
		}
		args, err := buildModuleRun(cfg, cat, e.ID)
		if err != nil {
			return fail("%v", err)
		}
		docker(40*time.Second, "rm", "-f", moduleContainer(e.ID))
		if _, err := docker(2*time.Minute, args...); err != nil {
			return fail("%v", err)
		}
		return ok(map[string]any{"name": moduleContainer(e.ID), "port": cat.HostPort(e.ID)})
	case "mod_stop":
		if _, err := docker(60*time.Second, "rm", "-f", moduleContainer(e.ID)); err != nil {
			return fail("%v", err)
		}
		return ok(nil)
	case "mod_logs":
		t := req.Tail
		if t < 1 || t > 200 {
			t = 30
		}
		out, err := dockerLogs(15*time.Second, "logs", "--tail", strconv.Itoa(t), moduleContainer(e.ID))
		if err != nil && out == "" {
			return fail("%v", err)
		}
		return ok(out)
	case "mod_rmi":
		// the image may be shared by other engines (same digest); the agent only asks when none of them is installed,
		// and a container that still uses it stops the removal
		if out, _ := docker(10*time.Second, "ps", "-a", "-q", "--filter", "label="+moduleLabel, "--filter", "ancestor="+e.ImageRef()); strings.TrimSpace(out) != "" {
			return fail("a module container still exists; stop it first")
		}
		if _, err := docker(5*time.Minute, "rmi", e.ImageRef()); err != nil {
			return fail("%v", err)
		}
		return ok(nil)
	}
	return fail("unknown operation")
}

// ourContainerIDs lists the full ids of the GPU containers that belong to the system: the module engines, and the ComfyUI
// lab container of step 0. The agent matches them against a GPU process's cgroup so the Owner Guard does not mistake
// them for the owner's own work.
func ourContainerIDs() []string {
	var ids []string
	for _, f := range [][]string{{"label=" + moduleLabel}, {"label=" + labelKey}, {"name=^wxch-comfy-lab$"}} {
		out, err := docker(10*time.Second, "ps", "-q", "--no-trunc", "--filter", f[0])
		if err != nil {
			continue
		}
		for _, id := range strings.Fields(out) {
			ids = append(ids, id)
		}
	}
	return ids
}