package main

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"math"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

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

// Jobs on the installed modules. The server sends a "job" message (capability + validated input); the agent starts the
// engine when it is not running, makes room in memory, fetches the input files from the server, runs the workflow (or
// asks Whisper), uploads the results to the server and answers with "job_done". The server keeps the queue, so the agent
// runs one job at a time and tells the server to try again (status 409) when it cannot take one: the owner is back, or
// it is already busy.

func randID() string {
	b := make([]byte, 6)
	rand.Read(b)
	return hex.EncodeToString(b)
}

// jobHost is where the agent reaches the server over HTTP (file transfer); set when the agent connects.
var jobHost = struct {
	sync.Mutex
	base, token string
}{}

func setJobHost(wsURL, token string) {
	u, err := url.Parse(wsURL)
	if err != nil {
		return
	}
	scheme := "https"
	if u.Scheme == "ws" {
		scheme = "http"
	}
	jobHost.Lock()
	jobHost.base, jobHost.token = scheme+"://"+u.Host, token
	jobHost.Unlock()
}

type jobStatusInfo struct {
	ID         string  `json:"id"`
	Capability string  `json:"capability"`
	Status     string  `json:"status"`
	Pct        float64 `json:"pct"`
	Started    int64   `json:"started"`
}

type jobRun struct {
	info   jobStatusInfo
	cancel context.CancelFunc
	reason string // why it was cancelled from this side
}

var jobs = struct {
	sync.Mutex
	running map[string]*jobRun
	lastUse map[string]time.Time // per engine: when it last finished a job
}{running: map[string]*jobRun{}, lastUse: map[string]time.Time{}}

func jobsBusy(capability string) bool {
	jobs.Lock()
	defer jobs.Unlock()
	for _, j := range jobs.running {
		if capability == "" || j.info.Capability == capability {
			return true
		}
	}
	return false
}

func jobsRunning() int {
	jobs.Lock()
	defer jobs.Unlock()
	return len(jobs.running)
}

func jobsOverview() []jobStatusInfo {
	jobs.Lock()
	defer jobs.Unlock()
	var out []jobStatusInfo
	for _, j := range jobs.running {
		out = append(out, j.info)
	}
	return out
}

// cancelAllJobs stops the running jobs from this side (the owner came back); the server queues them again.
func cancelAllJobs(reason string) {
	jobs.Lock()
	for _, j := range jobs.running {
		j.reason = reason
		j.cancel()
	}
	jobs.Unlock()
}

func cancelJob(id string) {
	jobs.Lock()
	if j := jobs.running[id]; j != nil {
		j.reason = "cancelled"
		j.cancel()
	}
	jobs.Unlock()
}

type jobParams struct {
	Input map[string]any `json:"input"`
}

type jobOutput struct {
	Name string `json:"name"`
	Size int64  `json:"size"`
	Mime string `json:"mime"`
}

type jobResult struct {
	Outputs []jobOutput `json:"outputs"`
	Text    string      `json:"text,omitempty"` // a short transcript, so the page can show it without a download
}

// handleJob serves one "job" message. It always ends with one job_done.
func handleJob(m Msg, send func(Msg)) {
	finish := func(status int, res *jobResult, err error) {
		out := Msg{T: "job_done", ID: m.ID, OK: err == nil, Status: status}
		if err != nil {
			out.Err = err.Error()
		}
		if res != nil {
			out.Data, _ = json.Marshal(res)
		}
		send(out)
	}
	if ok, why := guardAccepting(); !ok {
		finish(409, nil, fmt.Errorf("the machine is not taking jobs now: %s", why))
		return
	}
	if why := waitForLoad(context.Background(), jobLoadWait); why != "" {
		finish(409, nil, fmt.Errorf("%s", why))
		return
	}
	var p jobParams
	if err := json.Unmarshal(m.Params, &p); err != nil {
		finish(400, nil, fmt.Errorf("bad job"))
		return
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	run := &jobRun{info: jobStatusInfo{ID: m.ID, Capability: m.Method, Status: "starting", Started: time.Now().Unix()}, cancel: cancel}
	jobs.Lock()
	if len(jobs.running) > 0 {
		jobs.Unlock()
		finish(409, nil, fmt.Errorf("busy with another job"))
		return
	}
	jobs.running[m.ID] = run
	jobs.Unlock()
	report := func(status string, pct float64) {
		jobs.Lock()
		if status != "" {
			run.info.Status = status
		}
		if pct >= 0 {
			run.info.Pct = pct
		}
		st, pc := run.info.Status, run.info.Pct
		jobs.Unlock()
		b, _ := json.Marshal(map[string]any{"status": st, "pct": pc})
		send(Msg{T: "job_progress", ID: m.ID, Data: b})
	}
	res, engine, err := executeJob(ctx, loadNodeConfig(), m.ID, m.Method, p.Input, report)
	jobs.Lock()
	reason := run.reason
	delete(jobs.running, m.ID)
	if engine != "" {
		jobs.lastUse[engine] = time.Now()
	}
	jobs.Unlock()
	switch {
	case reason == "cancelled":
		finish(499, nil, fmt.Errorf("cancelled"))
	case reason != "":
		finish(409, nil, fmt.Errorf("stopped because %s", reason))
	case err != nil:
		finish(500, nil, err)
	default:
		finish(200, res, nil)
	}
}

func jobURL(path string, q url.Values) (string, string) {
	jobHost.Lock()
	defer jobHost.Unlock()
	return jobHost.base + path + "?" + q.Encode(), jobHost.token
}

var jobHTTP = &http.Client{Timeout: 30 * time.Minute}

func serverGet(ctx context.Context, id, name string) (io.ReadCloser, error) {
	u, tok := jobURL("/api/agent/job-file", url.Values{"job": {id}, "name": {name}, "dir": {"in"}})
	req, _ := http.NewRequestWithContext(ctx, "GET", u, nil)
	req.Header.Set("Authorization", "Bearer "+tok)
	resp, err := jobHTTP.Do(req)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != 200 {
		resp.Body.Close()
		return nil, fmt.Errorf("could not fetch %s from the server (%d)", name, resp.StatusCode)
	}
	return resp.Body, nil
}

func serverPut(ctx context.Context, id, name string, r io.Reader, size int64) error {
	u, tok := jobURL("/api/agent/job-file", url.Values{"job": {id}, "name": {name}, "dir": {"out"}})
	req, _ := http.NewRequestWithContext(ctx, "POST", u, r)
	req.Header.Set("Authorization", "Bearer "+tok)
	req.Header.Set("Content-Type", "application/octet-stream")
	if size >= 0 {
		req.ContentLength = size
	}
	resp, err := jobHTTP.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return &putStatusError{name, resp.StatusCode}
	}
	return nil
}

// putOutputWithRetry sends a finished result to the server. The work behind it can be minutes long, so a short outage of the
// server or the proxy in front of it (502/503/504, a dropped connection) is waited out instead of losing the result:
// the file is opened again for every try.
var putRetryWait = []time.Duration{5 * time.Second, 15 * time.Second, 45 * time.Second, 2 * time.Minute}

func putOutputWithRetry(ctx context.Context, id, name string, open func() (io.ReadCloser, int64, error)) (int64, error) {
	var last error
	for try := 0; ; try++ {
		body, size, err := open()
		if err != nil {
			return 0, err
		}
		err = serverPut(ctx, id, name, body, size)
		body.Close()
		if err == nil {
			return size, nil
		}
		last = err
		var se *putStatusError
		if errors.As(err, &se) && se.code != 502 && se.code != 503 && se.code != 504 {
			return 0, err // the server answered and refused: trying again will not change that
		}
		if try >= len(putRetryWait) {
			return 0, last
		}
		select {
		case <-ctx.Done():
			return 0, last
		case <-time.After(putRetryWait[try]):
		}
	}
}

type putStatusError struct {
	name string
	code int
}

func (e *putStatusError) Error() string {
	return fmt.Sprintf("the server did not accept %s (%d)", e.name, e.code)
}

var mimeByExt = map[string]string{".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".webp": "image/webp", ".gif": "image/gif",
	".mp4": "video/mp4", ".webm": "video/webm", ".mov": "video/quicktime", ".mkv": "video/x-matroska",
	".txt": "text/plain; charset=utf-8", ".srt": "application/x-subrip", ".vtt": "text/vtt", ".json": "application/json"}

func mimeOf(name string) string {
	if m := mimeByExt[strings.ToLower(filepath.Ext(name))]; m != "" {
		return m
	}
	return "application/octet-stream"
}

// engineHealthy waits until the engine answers.
func engineHealthy(ctx context.Context, base, health string, within time.Duration) error {
	deadline := time.Now().Add(within)
	hc := &http.Client{Timeout: 5 * time.Second}
	for {
		req, _ := http.NewRequestWithContext(ctx, "GET", base+health, nil)
		if resp, err := hc.Do(req); err == nil {
			resp.Body.Close()
			if resp.StatusCode == 200 {
				return nil
			}
		}
		if ctx.Err() != nil {
			return ctx.Err()
		}
		if time.Now().After(deadline) {
			return fmt.Errorf("the engine did not become ready in %v (see its log on the Modules page)", within)
		}
		time.Sleep(2 * time.Second)
	}
}

// engineBase is where an engine answers; a variable for tests.
var engineBase = func(cat *modules.Catalog, e modules.Engine) string {
	return fmt.Sprintf("http://127.0.0.1:%d", cat.HostPort(e.ID))
}

// engineWait is how long a cold engine may take (the first start installs Python packages): a variable for tests.
var engineWait = 30 * time.Minute

// ensureEngine starts the engine container when it is not running and waits until it answers.
func ensureEngine(ctx context.Context, cfg nodeConfig, cat *modules.Catalog, e modules.Engine, report func(string, float64)) (string, error) {
	rn, err := pickModRunner(cfg)
	if err != nil {
		return "", err
	}
	if ri, ok := rn.(runtimeInstaller); ok {
		if err := ri.EnsureRuntime(ctx, e, report); err != nil {
			return "", err
		}
	}
	running := false
	if st, err := rn.Status(); err == nil {
		for _, s := range st {
			if s.Engine == e.ID && s.State == "running" {
				running = true
			}
		}
	}
	base := engineBase(cat, e)
	if !running {
		report("starting the engine", -1)
		if err := rn.Start(e.ID); err != nil {
			return "", err
		}
	}
	return base, engineHealthy(ctx, base, e.Health, engineWait)
}

// makeRoom frees memory for a job by unloading language models when there is not enough, instead of refusing the job.
// (Models are loaded again by the next chat request.) An engine that ran a job a moment ago already holds its memory.
func makeRoom(ctx context.Context, cfg nodeConfig, engine string, need int64, report func(string, float64)) error {
	jobs.Lock()
	warm := time.Since(jobs.lastUse[engine]) < 10*time.Minute && !jobs.lastUse[engine].IsZero()
	jobs.Unlock()
	if warm || need <= 0 {
		return nil
	}
	reserve := uint64(cfg.ReserveGB) << 30
	enough := func() bool { _, avail := memInfo(); return avail >= uint64(need)+reserve }
	if enough() {
		return nil
	}
	report("making room in memory", -1)
	// only as many idle chat models as needed go, the least recently used first: the ones in use stay loaded
	if _, gone := thePool.freeUp(uint64(need)); len(gone) > 0 {
		logf("making room for a %s job: unloaded %s", engine, strings.Join(gone, ", "))
	}
	for i := 0; i < 30; i++ {
		if enough() {
			return nil
		}
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(2 * time.Second):
		}
	}
	_, avail := memInfo()
	return fmt.Errorf("not enough free memory: the job needs about %.0f GB (+%d GB kept for the owner), %.0f GB is free", float64(need)/1e9, cfg.ReserveGB, float64(avail)/1e9)
}

func executeJob(ctx context.Context, cfg nodeConfig, id, capID string, input map[string]any, report func(string, float64)) (*jobResult, string, error) {
	cat, err := loadCatalog()
	if err != nil {
		return nil, "", err
	}
	cp, ok := cat.Capability(capID)
	if !ok {
		return nil, "", fmt.Errorf("unknown capability %q", capID)
	}
	e, _ := cat.Engine(cp.Engine)
	if _, ok := readModuleState(cfg).Installed[capID]; !ok {
		return nil, "", fmt.Errorf("%s is not installed on this machine", capID)
	}
	in, err := cp.Validate(input)
	if err != nil {
		return nil, "", err
	}
	op := cp.Operation(in)
	base, err := ensureEngine(ctx, cfg, cat, e, report)
	if err != nil {
		return nil, e.ID, err
	}
	if err := makeRoom(ctx, cfg, e.ID, memNeed(cp, in), report); err != nil {
		return nil, e.ID, err
	}
	if e.ID == "whisper" {
		res, err := runWhisper(ctx, base, id, in, report)
		return res, e.ID, err
	}
	res, err := runComfy(ctx, cfg, cp, e, base, id, op, in, report)
	if err != nil && ctx.Err() == nil {
		err = withEngineLog(cfg, e.ID, err)
	}
	return res, e.ID, err
}

const maxLogInError = 4000

// withEngineLog adds the last lines of the engine's own log to the error of a failed job. The log holds the traceback that
// says where a CUDA or memory problem happened, and it is gone when the engine is restarted, so it is kept with the job.
func withEngineLog(cfg nodeConfig, engine string, err error) error {
	rn, rerr := pickModRunner(cfg)
	if rerr != nil {
		return err
	}
	logs := strings.TrimSpace(rn.Logs(engine, 80))
	if logs == "" {
		return err
	}
	if len(logs) > maxLogInError {
		logs = "…" + logs[len(logs)-maxLogInError:]
	}
	return fmt.Errorf("%v\n\n--- engine log (last lines) ---\n%s", err, logs)
}

// memNeed is the memory a job needs. The figure in the catalog is for a short clip; a longer video has more frames to hold
// (the latents and the decoding grow with the length). The growth per second is a rough guess until it has been measured.
const videoMemPerExtraSecond = int64(1.5e9)

func asFloat(v any) (float64, bool) {
	switch x := v.(type) {
	case float64:
		return x, true
	case int64:
		return float64(x), true
	case int:
		return float64(x), true
	}
	return 0, false
}

// storySeconds lists the lengths of the scenes of a story input.
func storySeconds(scenes any) []float64 {
	var out []float64
	list, _ := scenes.([]any)
	for _, e := range list {
		if m, ok := e.(map[string]any); ok {
			if s, ok := asFloat(m["seconds"]); ok {
				out = append(out, s)
			}
		}
	}
	return out
}

func memNeed(cp modules.Capability, in map[string]any) int64 {
	need := cp.MemBytes
	if cp.ID == "video.generate" {
		if cp.Operation(in) == "story" {
			// one scene at a time is on the GPU; what grows is the joined frames, kept in memory until the video is saved
			var longest float64
			for _, s := range storySeconds(in["scenes"]) {
				longest = math.Max(longest, s)
			}
			if longest > 3 {
				need += int64((longest - 3) * float64(videoMemPerExtraSecond))
			}
			size, _ := in["size"].(string)
			if size == "" {
				size = "832x480"
			}
			need += modules.StoryMemBytes(in["scenes"], size)
		} else if sec, ok := asFloat(in["seconds"]); ok && sec > 3 {
			need += int64((sec - 3) * float64(videoMemPerExtraSecond))
		}
	}
	return need
}

func runComfy(ctx context.Context, cfg nodeConfig, cp modules.Capability, e modules.Engine, base, id, op string, in map[string]any, report func(string, float64)) (*jobResult, error) {
	c := newComfy(base)
	root := filepath.Join(engineDir(cfg, e), filepath.Base(e.Workdir))
	defer os.RemoveAll(filepath.Join(root, "input", "wxch", id))
	defer os.RemoveAll(filepath.Join(root, "output", "wxch", id))
	vals, err := cp.Values(in)
	if err != nil {
		return nil, err
	}
	for _, f := range cp.FileFields(op) {
		name, _ := in[f.Name].(string)
		if name == "" {
			continue
		}
		report("sending "+name+" to the engine", -1)
		body, err := serverGet(ctx, id, name)
		if err != nil {
			return nil, err
		}
		up, err := c.Upload(ctx, "wxch/"+id, name, body)
		body.Close()
		if err != nil {
			return nil, err
		}
		vals[f.Name] = up
	}
	vals["prefix"] = "wxch/" + id + "/out"
	file, err := cp.WorkflowFor(op)
	if err != nil {
		return nil, err
	}
	var wf map[string]any
	if file == modules.BuiltinStory { // several scenes: the graph is built in code
		if wf, err = modules.BuildStory(vals); err != nil {
			return nil, err
		}
	} else {
		raw, err := modules.Workflow(file)
		if err != nil {
			return nil, err
		}
		if wf, err = modules.Instantiate(raw, vals); err != nil {
			return nil, err
		}
	}
	report("running", 0)
	outs, err := c.Run(ctx, wf, func(p float64) { report("", p) })
	if err != nil {
		return nil, err
	}
	if len(outs) == 0 {
		return nil, fmt.Errorf("the engine finished but saved nothing")
	}
	report("saving the result", 1)
	res := &jobResult{}
	for _, o := range outs {
		size, err := putOutputWithRetry(ctx, id, o.Filename, func() (io.ReadCloser, int64, error) { return c.Fetch(ctx, o) })
		if err != nil {
			return nil, err
		}
		res.Outputs = append(res.Outputs, jobOutput{Name: o.Filename, Size: size, Mime: mimeOf(o.Filename)})
	}
	return res, nil
}

// handleWarm starts the engine behind a capability ahead of time, because somebody opened the page that will use it. It
// starts only the engine (no model is loaded, nothing is unloaded) and does nothing when the machine is not taking work,
// the capability is not installed here, or the engine is up already (then it only counts as used, so the idle clock
// restarts). Asked again and again while the page is open, it keeps the engine for as long as somebody is there.
func handleWarm(m Msg) {
	cfg := loadNodeConfig()
	cat, err := loadCatalog()
	if err != nil {
		return
	}
	cp, ok := cat.Capability(m.Method)
	if !ok {
		return
	}
	if _, installed := readModuleState(cfg).Installed[cp.ID]; !installed {
		return
	}
	if accepting, _ := guardAccepting(); !accepting {
		return
	}
	e, _ := cat.Engine(cp.Engine)
	touchEngine(e.ID)
	maybeStartEngine(cfg, cat, e)
}

// ---- idle stop ----

// moduleIdleLoop stops an engine that has had no job for a while, so its memory goes back to the language models.
func moduleIdleLoop(stop <-chan struct{}) {
	t := time.NewTicker(30 * time.Second)
	defer t.Stop()
	for {
		select {
		case <-stop:
			return
		case <-t.C:
		}
		cfg := loadNodeConfig()
		if len(readModuleState(cfg).Installed) == 0 || jobsRunning() > 0 || tunnelsOpen() {
			continue
		}
		if cfg.ModuleIdleS < 0 { // module_idle_s: -1 = never stop an engine automatically
			continue
		}
		idle := time.Duration(cfg.ModuleIdleS) * time.Second
		if idle == 0 {
			idle = 10 * time.Minute
		}
		rn, err := pickModRunner(cfg)
		if err != nil {
			continue
		}
		st, err := rn.Status()
		if err != nil {
			continue
		}
		for _, e := range enginesToStop(st, idle, time.Now(), readModuleState(cfg).Pinned) {
			if _, already := stopping.LoadOrStore(e, true); already {
				continue
			}
			go func(e string) {
				defer stopping.Delete(e)
				logf("module engine %s idle for %v: stopping it", e, idle)
				rn.Stop(e)
			}(e)
		}
	}
}

var stopping sync.Map // engines being stopped right now

// enginesToStop says which running engines have had no use for longer than idle. An engine that is still starting up is
// not idle: nobody can use it yet, and stopping it would throw away the first start, the slowest one (the Python packages
// are installed then). Its idle time counts from the moment it is ready.
func enginesToStop(st []runnerModule, idle time.Duration, now time.Time, pinned map[string]bool) []string {
	jobs.Lock()
	defer jobs.Unlock()
	var out []string
	for _, s := range st {
		if s.State != "running" || pinned[s.Engine] { // an engine started by hand stays until it is stopped by hand
			continue
		}
		if _, starting := startingEngine.Load(s.Engine); starting {
			jobs.lastUse[s.Engine] = now
			continue
		}
		last := jobs.lastUse[s.Engine]
		if last.IsZero() {
			jobs.lastUse[s.Engine] = now // started by hand: count from now
			continue
		}
		if now.Sub(last) > idle {
			out = append(out, s.Engine)
		}
	}
	return out
}

// installedModuleIDs is what telemetry reports so the server knows which machine can do what.
func installedModuleIDs(cfg nodeConfig) []string {
	var ids []string
	for id := range readModuleState(cfg).Installed {
		ids = append(ids, id)
	}
	return ids
}