package main

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

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

// fakeComfy answers like ComfyUI for one workflow at a time.
type fakeComfy struct {
	mu          sync.Mutex
	prompt      map[string]any
	uploaded    []string
	hang        bool // never finish, to test cancelling
	interrupted bool
	fail        bool // answer with an execution error
	srv         *httptest.Server
}

func newFakeComfy(t *testing.T) *fakeComfy {
	f := &fakeComfy{}
	mux := http.NewServeMux()
	mux.HandleFunc("/system_stats", func(w http.ResponseWriter, r *http.Request) { w.Write([]byte("{}")) })
	mux.HandleFunc("/upload/image", func(w http.ResponseWriter, r *http.Request) {
		r.ParseMultipartForm(1 << 20)
		_, h, _ := r.FormFile("image")
		f.mu.Lock()
		f.uploaded = append(f.uploaded, r.FormValue("subfolder")+"/"+h.Filename)
		f.mu.Unlock()
		json.NewEncoder(w).Encode(map[string]string{"name": h.Filename, "subfolder": r.FormValue("subfolder"), "type": "input"})
	})
	mux.HandleFunc("/prompt", func(w http.ResponseWriter, r *http.Request) {
		var b struct{ Prompt map[string]any }
		json.NewDecoder(r.Body).Decode(&b)
		f.mu.Lock()
		f.prompt = b.Prompt
		f.mu.Unlock()
		json.NewEncoder(w).Encode(map[string]string{"prompt_id": "p1"})
	})
	mux.HandleFunc("/history/p1", func(w http.ResponseWriter, r *http.Request) {
		f.mu.Lock()
		hang := f.hang
		f.mu.Unlock()
		fail := f.fail
		if hang {
			w.Write([]byte("{}"))
			return
		}
		if fail {
			fmt.Fprint(w, `{"p1":{"status":{"status_str":"error","completed":false,"messages":[["execution_error",{"node_type":"KSampler","exception_message":"CUDA error: operation not permitted"}]]},"outputs":{}}}`)
			return
		}
		fmt.Fprint(w, `{"p1":{"status":{"status_str":"success","completed":true},"outputs":{"9":{"images":[{"filename":"out_00001_.png","subfolder":"wxch/J1","type":"output"}]}}}}`)
	})
	mux.HandleFunc("/history", func(w http.ResponseWriter, r *http.Request) {})
	mux.HandleFunc("/queue", func(w http.ResponseWriter, r *http.Request) {})
	mux.HandleFunc("/interrupt", func(w http.ResponseWriter, r *http.Request) { f.mu.Lock(); f.interrupted = true; f.mu.Unlock() })
	mux.HandleFunc("/view", func(w http.ResponseWriter, r *http.Request) { w.Write([]byte("PNGDATA")) })
	f.srv = httptest.NewServer(mux)
	t.Cleanup(f.srv.Close)
	return f
}

// fakeControl is the server's file endpoint for jobs.
type fakeControl struct {
	mu  sync.Mutex
	out map[string]string
	srv *httptest.Server
}

func newFakeControl(t *testing.T) *fakeControl {
	c := &fakeControl{out: map[string]string{}}
	c.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Authorization") != "Bearer tok" {
			http.Error(w, "no", 401)
			return
		}
		q := r.URL.Query()
		if r.Method == "GET" {
			w.Write([]byte("INPUT-" + q.Get("name")))
			return
		}
		b, _ := io.ReadAll(r.Body)
		c.mu.Lock()
		c.out[q.Get("name")] = string(b)
		c.mu.Unlock()
	}))
	t.Cleanup(c.srv.Close)
	setJobHost(strings.Replace(c.srv.URL, "http", "ws", 1)+"/ws/agent", "tok")
	return c
}

func jobSetup(t *testing.T) (nodeConfig, *fakeComfy, *fakeControl) {
	cfg := nodeConfig{ModulesDir: filepath.Join(t.TempDir(), "modules"), ReserveGB: 8}
	os.MkdirAll(filepath.Join(cfg.ModulesDir, "comfyui", "ComfyUI"), 0755)
	st := moduleState{Installed: map[string]string{"image.generate": "x", "image.edit": "x"}}
	writeModuleState(cfg, st)
	fc := newFakeComfy(t)
	ctl := newFakeControl(t)
	rn := &fakeModRunner{image: map[string]bool{"comfyui": true}, running: map[string]bool{"comfyui": true}}
	modRunnerOverride = rn
	oldBase := engineBase
	engineBase = func(*modules.Catalog, modules.Engine) string { return fc.srv.URL }
	jobs.Lock()
	jobs.lastUse["comfyui"] = time.Now() // a warm engine: no memory juggling in the test
	jobs.running = map[string]*jobRun{}
	jobs.Unlock()
	guardState.Lock()
	guardState.g = nil
	guardState.Unlock()
	t.Cleanup(func() { modRunnerOverride = nil; engineBase = oldBase })
	return cfg, fc, ctl
}

func runJobMsg(t *testing.T, cfg nodeConfig, id, capID string, input map[string]any) []Msg {
	_ = cfg
	var mu sync.Mutex
	var got []Msg
	params, _ := json.Marshal(map[string]any{"input": input})
	handleJob(Msg{T: "job", ID: id, Method: capID, Params: params}, func(m Msg) { mu.Lock(); got = append(got, m); mu.Unlock() })
	return got
}

func TestJobRunsAWorkflowAndReturnsTheFiles(t *testing.T) {
	cfgOverride = &nodeConfig{ModulesDir: "", ReserveGB: 8}
	cfg, fc, ctl := jobSetup(t)
	cfgOverride = &cfg
	t.Cleanup(func() { cfgOverride = nil })
	msgs := runJobMsg(t, cfg, "J1", "image.edit", map[string]any{"image": "in.png", "prompt": "make it blue"})
	done := msgs[len(msgs)-1]
	if done.T != "job_done" || !done.OK || done.Status != 200 {
		t.Fatalf("job_done = %+v", done)
	}
	var res jobResult
	json.Unmarshal(done.Data, &res)
	if len(res.Outputs) != 1 || res.Outputs[0].Name != "out_00001_.png" || res.Outputs[0].Mime != "image/png" {
		t.Fatalf("outputs = %+v", res)
	}
	if ctl.out["out_00001_.png"] != "PNGDATA" {
		t.Errorf("the result was not sent to the server: %v", ctl.out)
	}
	if len(fc.uploaded) != 1 || fc.uploaded[0] != "wxch/J1/in.png" {
		t.Errorf("input upload = %v", fc.uploaded)
	}
	wf, _ := json.Marshal(fc.prompt)
	for _, want := range []string{"make it blue", "wxch/J1/in.png", "wxch/J1/out"} {
		if !strings.Contains(string(wf), want) {
			t.Errorf("the workflow sent to the engine lacks %q", want)
		}
	}
	if strings.Contains(string(wf), `"@`) {
		t.Error("a placeholder was left in the workflow")
	}
	var sawProgress bool
	for _, m := range msgs {
		if m.T == "job_progress" {
			sawProgress = true
		}
	}
	if !sawProgress {
		t.Error("no progress was reported")
	}
	if jobsRunning() != 0 {
		t.Error("the job is still listed as running")
	}
}

func TestJobRefusedWhenNotInstalledOrInvalid(t *testing.T) {
	cfg, _, _ := jobSetup(t)
	cfgOverride = &cfg
	t.Cleanup(func() { cfgOverride = nil })
	if m := runJobMsg(t, cfg, "J2", "video.generate", map[string]any{"prompt": "x"}); m[len(m)-1].OK {
		t.Error("a capability that is not installed ran")
	}
	if m := runJobMsg(t, cfg, "J3", "image.generate", map[string]any{}); m[len(m)-1].OK {
		t.Error("a job without a prompt ran")
	}
}

func TestJobBouncesWhenTheOwnerIsBackOrBusy(t *testing.T) {
	cfg, _, _ := jobSetup(t)
	cfgOverride = &cfg
	t.Cleanup(func() { cfgOverride = nil })
	guardState.Lock()
	guardState.g = newOwnerGuard(defaultGuardConfig())
	guardState.last = guardDecision{Mode: modeOwner, Reason: "owner is using the machine", StopAccepting: true}
	guardState.Unlock()
	m := runJobMsg(t, cfg, "J4", "image.generate", map[string]any{"prompt": "x"})
	if d := m[len(m)-1]; d.OK || d.Status != 409 {
		t.Fatalf("while the owner is here the server must be told to try later: %+v", d)
	}
	guardState.Lock()
	guardState.g = nil
	guardState.Unlock()
	jobs.Lock()
	jobs.running["other"] = &jobRun{}
	jobs.Unlock()
	m = runJobMsg(t, cfg, "J5", "image.generate", map[string]any{"prompt": "x"})
	jobs.Lock()
	delete(jobs.running, "other")
	jobs.Unlock()
	if d := m[len(m)-1]; d.Status != 409 {
		t.Fatalf("a second job at once must be bounced: %+v", d)
	}
}

func TestJobCancelAndOwnerReturnStopTheEngine(t *testing.T) {
	cfg, fc, _ := jobSetup(t)
	cfgOverride = &cfg
	t.Cleanup(func() { cfgOverride = nil })
	for _, tc := range []struct {
		name   string
		stop   func(id string)
		status int
	}{
		{"cancel", func(id string) { cancelJob(id) }, 499},
		{"owner returns", func(string) { cancelAllJobs("the owner is using the machine") }, 409},
	} {
		fc.mu.Lock()
		fc.hang, fc.interrupted = true, false
		fc.mu.Unlock()
		res := make(chan []Msg)
		go func() { res <- runJobMsg(t, cfg, "JC", "image.generate", map[string]any{"prompt": "x"}) }()
		for i := 0; i < 100 && !jobsBusy("image.generate"); i++ {
			time.Sleep(20 * time.Millisecond)
		}
		time.Sleep(1500 * time.Millisecond) // the job is waiting on the engine
		tc.stop("JC")
		select {
		case m := <-res:
			if d := m[len(m)-1]; d.OK || d.Status != tc.status {
				t.Errorf("%s: %+v", tc.name, d)
			}
		case <-time.After(10 * time.Second):
			t.Fatalf("%s: the job did not stop", tc.name)
		}
		fc.mu.Lock()
		if !fc.interrupted {
			t.Errorf("%s: the engine was not told to stop", tc.name)
		}
		fc.mu.Unlock()
	}
}

func TestSubtitles(t *testing.T) {
	segs := []whisperSeg{{0, 2.5, "hello"}, {3661.007, 3662, "world"}}
	srt := toSRT(segs)
	if !strings.Contains(srt, "1\n00:00:00,000 --> 00:00:02,500\nhello") || !strings.Contains(srt, "2\n01:01:01,007 --> 01:01:02,000\nworld") {
		t.Errorf("srt = %q", srt)
	}
	if !strings.HasPrefix(toVTT(segs), "WEBVTT\n\n00:00:00.000 --> 00:00:02.500\nhello") {
		t.Errorf("vtt = %q", toVTT(segs))
	}
}

func TestAnEngineThatIsStillStartingIsNeverStoppedAsIdle(t *testing.T) {
	now := time.Now()
	st := []runnerModule{{Engine: "comfyui", State: "running"}, {Engine: "whisper", State: "running"}, {Engine: "other", State: "exited"}}
	idle := 10 * time.Minute
	jobs.Lock()
	jobs.lastUse = map[string]time.Time{"comfyui": now.Add(-time.Hour), "whisper": now.Add(-time.Hour), "other": now.Add(-time.Hour)}
	jobs.Unlock()
	startingEngine.Delete("comfyui")

	got := enginesToStop(st, idle, now, nil)
	if len(got) != 2 {
		t.Fatalf("engines idle for an hour must be stopped (but not one that is not running): %v", got)
	}
	// comfyui is still being started: the first start is the slow one and must not be cut short
	startingEngine.Store("comfyui", true)
	t.Cleanup(func() { startingEngine.Delete("comfyui") })
	got = enginesToStop(st, idle, now, nil)
	if len(got) != 1 || got[0] != "whisper" {
		t.Fatalf("an engine that is starting up is not idle: %v", got)
	}
	// and the clock counts from when it became ready, not from when it started
	startingEngine.Delete("comfyui")
	if got := enginesToStop(st, idle, now.Add(5*time.Minute), nil); len(got) != 1 || got[0] != "whisper" {
		t.Errorf("five minutes after it became ready it is still in use: %v", got)
	}
	if got := enginesToStop(st, idle, now.Add(11*time.Minute), nil); len(got) != 2 {
		t.Errorf("eleven minutes after it became ready it is idle: %v", got)
	}
	// an engine started by hand (no use recorded) is counted from now
	jobs.Lock()
	jobs.lastUse = map[string]time.Time{}
	jobs.Unlock()
	if got := enginesToStop(st, idle, now, nil); len(got) != 0 {
		t.Errorf("an engine with no recorded use starts its clock now: %v", got)
	}
}

func TestAnEngineStartedByHandStaysUntilStoppedByHand(t *testing.T) {
	now := time.Now()
	st := []runnerModule{{Engine: "comfyui", State: "running"}, {Engine: "whisper", State: "running"}}
	jobs.Lock()
	jobs.lastUse = map[string]time.Time{"comfyui": now.Add(-5 * time.Hour), "whisper": now.Add(-5 * time.Hour)}
	jobs.Unlock()
	startingEngine.Delete("comfyui")
	got := enginesToStop(st, 10*time.Minute, now, map[string]bool{"comfyui": true})
	if len(got) != 1 || got[0] != "whisper" {
		t.Fatalf("only the engine that was not started by hand may be stopped: %v", got)
	}
}

func TestPinningFollowsTheStartAndStopButtons(t *testing.T) {
	cfg := nodeConfig{ModulesDir: t.TempDir()}
	writeModuleState(cfg, moduleState{Installed: map[string]string{"image.generate": "x"}})
	setPinned(cfg, "comfyui", true)
	if !readModuleState(cfg).Pinned["comfyui"] {
		t.Fatal("start by hand must pin the engine")
	}
	if len(readModuleState(cfg).Installed) != 1 {
		t.Error("pinning must not touch what is installed")
	}
	setPinned(cfg, "comfyui", false)
	if readModuleState(cfg).Pinned["comfyui"] {
		t.Error("stop by hand must unpin it")
	}
}

func TestWarmStartsOnlyTheEngineAndOnlyWhenItMakesSense(t *testing.T) {
	cfg, rn, _ := moduleSetup(t) // fake runner, a.one/a.two/s.one in a test catalog
	_ = rn
	// use the real catalog: image.generate runs on comfyui
	loadCatalogOld := loadCatalog
	loadCatalog = modules.Load
	t.Cleanup(func() { loadCatalog = loadCatalogOld })
	cfgOverride = &cfg
	t.Cleanup(func() { cfgOverride = nil })
	fake := &fakeModRunner{image: map[string]bool{"comfyui": true}}
	modRunnerOverride = fake
	guardState.Lock()
	guardState.g = nil
	guardState.Unlock()
	started := func() bool { fake.mu.Lock(); defer fake.mu.Unlock(); return fake.running["comfyui"] }
	reset := func() {
		fake.mu.Lock()
		fake.running = nil
		fake.mu.Unlock()
		lastEngineStart.Delete("comfyui")
		for i := 0; i < 300; i++ { // an earlier start may still be waiting for its health check
			if _, busy := startingEngine.Load("comfyui"); !busy {
				break
			}
			time.Sleep(20 * time.Millisecond)
		}
	}
	oldWait := engineWait
	engineWait = 500 * time.Millisecond
	t.Cleanup(func() { engineWait = oldWait })

	// not installed: nothing happens
	reset()
	handleWarm(Msg{T: "warm", Method: "image.generate"})
	time.Sleep(150 * time.Millisecond)
	if started() {
		t.Error("an engine of a capability that is not installed was started")
	}
	writeModuleState(cfg, moduleState{Installed: map[string]string{"image.generate": "x"}})
	// unknown capability
	handleWarm(Msg{T: "warm", Method: "nope"})
	// the machine is not taking work (the owner is here)
	guardState.Lock()
	guardState.g = newOwnerGuard(defaultGuardConfig())
	guardState.last = guardDecision{Mode: modeOwner, StopAccepting: true}
	guardState.Unlock()
	handleWarm(Msg{T: "warm", Method: "image.generate"})
	time.Sleep(150 * time.Millisecond)
	if started() {
		t.Error("the engine was started while the owner is using the machine")
	}
	guardState.Lock()
	guardState.g = nil
	guardState.Unlock()
	// the normal case
	reset()
	handleWarm(Msg{T: "warm", Method: "image.generate"})
	for i := 0; i < 100 && !started(); i++ {
		time.Sleep(20 * time.Millisecond)
	}
	if !started() {
		t.Fatal("the engine was not started")
	}
	// asked again while it runs: nothing is started again, but it counts as used
	jobs.Lock()
	jobs.lastUse["comfyui"] = time.Now().Add(-time.Hour)
	jobs.Unlock()
	handleWarm(Msg{T: "warm", Method: "image.generate"})
	jobs.Lock()
	recent := time.Since(jobs.lastUse["comfyui"]) < time.Minute
	jobs.Unlock()
	if !recent {
		t.Error("a repeated request must restart the idle clock")
	}
	time.Sleep(700 * time.Millisecond)
}

func TestAFailedJobCarriesTheEngineLog(t *testing.T) {
	cfg, fc, _ := jobSetup(t)
	cfgOverride = &cfg
	t.Cleanup(func() { cfgOverride = nil })
	fc.mu.Lock()
	fc.fail = true
	fc.mu.Unlock()
	rn := modRunnerOverride.(*fakeModRunner)
	rn.logs = "Traceback (most recent call last):" + string(rune(10)) + "  File \"comfy/model_management.py\", line 1, in load" + string(rune(10)) + "RuntimeError: CUDA error: operation not permitted"
	msgs := runJobMsg(t, cfg, "JL", "image.generate", map[string]any{"prompt": "x"})
	done := msgs[len(msgs)-1]
	if done.OK || done.Status != 500 {
		t.Fatalf("%+v", done)
	}
	if !strings.Contains(done.Err, "KSampler: CUDA error: operation not permitted") || !strings.Contains(done.Err, "engine log") || !strings.Contains(done.Err, "model_management.py") {
		t.Errorf("the error must say what failed and carry the traceback: %q", done.Err)
	}
	// a very long log is cut to its end
	rn.logs = strings.Repeat("x", 9000) + "THE-END"
	msgs = runJobMsg(t, cfg, "JM", "image.generate", map[string]any{"prompt": "x"})
	if e := msgs[len(msgs)-1].Err; len(e) > 4600 || !strings.HasSuffix(e, "THE-END") {
		t.Errorf("the log must be cut to its last part: %d bytes", len(e))
	}
	// no log, no extra text
	rn.logs = ""
	msgs = runJobMsg(t, cfg, "JN", "image.generate", map[string]any{"prompt": "x"})
	if strings.Contains(msgs[len(msgs)-1].Err, "engine log") {
		t.Error("an empty log must not add a heading")
	}
}

func TestAJobNeedsMoreMemoryTheLongerTheClip(t *testing.T) {
	c, err := modules.Load()
	if err != nil {
		t.Fatal(err)
	}
	vg, _ := c.Capability("video.generate")
	short := memNeed(vg, map[string]any{"seconds": float64(3)})
	long := memNeed(vg, map[string]any{"seconds": float64(30)})
	if short != vg.MemBytes || long <= short || long > short+50e9 {
		t.Errorf("3 s = %d, 30 s = %d", short, long)
	}
	if memNeed(vg, map[string]any{}) != vg.MemBytes {
		t.Error("no length given: the catalog figure")
	}
	ig, _ := c.Capability("image.generate")
	if memNeed(ig, map[string]any{"seconds": float64(20)}) != ig.MemBytes {
		t.Error("only video grows with its length")
	}
}

func TestAStoryJobSendsOneGraphWithAllTheScenes(t *testing.T) {
	cfg, fc, ctl := jobSetup(t)
	cfgOverride = &cfg
	t.Cleanup(func() { cfgOverride = nil })
	writeModuleState(cfg, moduleState{Installed: map[string]string{"video.generate": "x"}})
	msgs := runJobMsg(t, cfg, "J1", "video.generate", map[string]any{
		"operation": "story", "style": "soft light",
		"scenes": []any{
			map[string]any{"prompt": "a pink hallway", "seconds": 4},
			map[string]any{"prompt": "a floating bedroom", "seconds": 5.0},
			map[string]any{"prompt": "a playground"},
		},
	})
	done := msgs[len(msgs)-1]
	if done.T != "job_done" || !done.OK || done.Status != 200 {
		t.Fatalf("job_done = %+v %s", done, done.Err)
	}
	if ctl.out["out_00001_.png"] == "" {
		t.Errorf("the result was not sent to the server: %v", ctl.out)
	}
	fc.mu.Lock()
	wf := fc.prompt
	fc.mu.Unlock()
	classes := map[string]int{}
	for _, n := range wf {
		classes[n.(map[string]any)["class_type"].(string)]++
	}
	if classes["KSampler"] != 3 || classes["ImageBatch"] != 2 || classes["SaveVideo"] != 1 || classes["VAEDecodeTiled"] != 3 {
		t.Errorf("graph: %v", classes)
	}
	b, _ := json.Marshal(wf)
	for _, want := range []string{"a pink hallway", "a floating bedroom", "a playground", "soft light", "wxch/J1/out"} {
		if !strings.Contains(string(b), want) {
			t.Errorf("the graph lacks %q", want)
		}
	}
	if strings.Contains(string(b), `"@`) {
		t.Error("a placeholder was left in the graph")
	}
}

func TestAStoryNeedsMemoryForItsJoinedFrames(t *testing.T) {
	c, err := modules.Load()
	if err != nil {
		t.Fatal(err)
	}
	vg, _ := c.Capability("video.generate")
	one := map[string]any{"operation": "story", "size": "832x480", "scenes": []any{map[string]any{"prompt": "x", "seconds": 3.0}}}
	many := map[string]any{"operation": "story", "size": "832x480", "scenes": []any{
		map[string]any{"prompt": "x", "seconds": 8.0}, map[string]any{"prompt": "y", "seconds": 8.0}, map[string]any{"prompt": "z", "seconds": 8.0}}}
	if memNeed(vg, many) <= memNeed(vg, one) || memNeed(vg, one) < vg.MemBytes {
		t.Errorf("one scene %d, three long scenes %d", memNeed(vg, one), memNeed(vg, many))
	}
	big := map[string]any{"operation": "story", "size": "1280x704", "scenes": many["scenes"]}
	if memNeed(vg, big) <= memNeed(vg, many) {
		t.Error("a bigger picture needs more memory")
	}
}