package main

import (
	"bytes"
	"encoding/json"
	"io"
	"mime/multipart"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

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

type studioFixture struct {
	s      *Server
	mux    http.Handler
	cookie *http.Cookie
	mu     sync.Mutex
	sent   []Msg
}

func newStudioFixture(t *testing.T) *studioFixture {
	cat, err := modules.Load()
	if err != nil {
		t.Fatal(err)
	}
	st := newSessionStore("")
	s := &Server{studioSess: st.set("studio"), cfg: Config{AgentTokens: map[string]string{"pgx": "tok-pgx", "other": "tok-other"}},
		hub:  &Hub{agents: map[string]*conn{}, uis: map[*conn]bool{}, sessions: map[string]*session{}, rpcs: map[string]*rpcEntry{}, stats: map[string]*statBuf{}, infers: map[string]*inferJob{}, inflight: map[string]int{}},
		jobs: newJobStore(t.TempDir(), cat)}
	f := &studioFixture{s: s, mux: s.studioMux(os.DirFS("../../web/static"))}
	s.studioSess.add("session-token", time.Hour)
	f.cookie = &http.Cookie{Name: "wxch_studio", Value: "session-token"}
	return f
}

// machine adds an online machine that reports these modules.
func (f *studioFixture) machine(name, mode string, mods ...string) {
	c := &conn{sink: func(m Msg) { f.mu.Lock(); f.sent = append(f.sent, m); f.mu.Unlock() }}
	f.s.hub.mu.Lock()
	f.s.hub.agents[name] = c
	f.s.hub.mu.Unlock()
	b, _ := json.Marshal(map[string]any{"ts": time.Now().Unix(), "mode": mode, "modules": mods, "mem_avail": 90e9, "mem_total": 120e9})
	f.s.hub.recordStats(name, b)
}

func (f *studioFixture) do(method, path string, body io.Reader, ctype string, authed bool) *httptest.ResponseRecorder {
	r := httptest.NewRequest(method, path, body)
	if ctype != "" {
		r.Header.Set("Content-Type", ctype)
	}
	if authed {
		r.AddCookie(f.cookie)
	}
	w := httptest.NewRecorder()
	f.mux.ServeHTTP(w, r)
	return w
}

func form(fields map[string]string, files map[string][2]string) (*bytes.Buffer, string) {
	var b bytes.Buffer
	mw := multipart.NewWriter(&b)
	mw.WriteField("capability", fields["capability"])
	for f, nc := range files {
		fw, _ := mw.CreateFormFile(f, nc[0])
		fw.Write([]byte(nc[1]))
	}
	mw.WriteField("input", fields["input"])
	mw.Close()
	return &b, mw.FormDataContentType()
}

func (f *studioFixture) submit(capID, input string, files map[string][2]string) *httptest.ResponseRecorder {
	b, ct := form(map[string]string{"capability": capID, "input": input}, files)
	return f.do("POST", "/studio/api/jobs", b, ct, true)
}

func (f *studioFixture) agentReq(token, method, query string, body io.Reader) *httptest.ResponseRecorder {
	r := httptest.NewRequest(method, "/api/agent/job-file?"+query, body)
	r.Header.Set("Authorization", "Bearer "+token)
	w := httptest.NewRecorder()
	f.s.agentJobFile(w, r)
	return w
}

func decodeJob(t *testing.T, w *httptest.ResponseRecorder) Job {
	var j Job
	if err := json.Unmarshal(w.Body.Bytes(), &j); err != nil {
		t.Fatalf("%v: %s", err, w.Body.String())
	}
	return j
}

func TestStudioNeedsLogin(t *testing.T) {
	f := newStudioFixture(t)
	for _, p := range []string{"/studio/api/capabilities", "/studio/api/jobs"} {
		if w := f.do("GET", p, nil, "", false); w.Code != 401 {
			t.Errorf("%s without a session: %d", p, w.Code)
		}
	}
	if w := f.do("GET", "/", nil, "", false); w.Code != 200 || !strings.Contains(w.Body.String(), "wxch studio") {
		t.Errorf("the page itself must load to show the login: %d", w.Code)
	}
}

func TestOnlyWhatMachinesInstalledIsOffered(t *testing.T) {
	f := newStudioFixture(t)
	if w := f.submit("image.generate", `{"prompt":"x"}`, nil); w.Code != 503 {
		t.Fatalf("with no machine the job must be refused, got %d", w.Code)
	}
	f.machine("pgx", "available", "image.generate")
	var caps struct{ Capabilities []studioCap }
	json.Unmarshal(f.do("GET", "/studio/api/capabilities", nil, "", true).Body.Bytes(), &caps)
	got := map[string]studioCap{}
	for _, c := range caps.Capabilities {
		got[c.ID] = c
	}
	if len(got) != 5 || !got["image.generate"].Offered || !got["image.generate"].Available {
		t.Fatalf("image.generate must be offered and free: %+v", got["image.generate"])
	}
	if got["video.generate"].Offered {
		t.Error("a capability no machine installed must not be offered")
	}
	if w := f.submit("video.generate", `{"prompt":"x"}`, nil); w.Code != 503 {
		t.Errorf("a job for a capability nobody installed: %d", w.Code)
	}
}

func TestAJobGoesToTheMachineAndBack(t *testing.T) {
	f := newStudioFixture(t)
	f.machine("pgx", "available", "image.generate")
	w := f.submit("image.generate", `{"prompt":"a cat","count":2}`, nil)
	if w.Code != 200 {
		t.Fatalf("%d %s", w.Code, w.Body.String())
	}
	j := decodeJob(t, w)
	if j.Status != "queued" {
		t.Fatalf("status %s", j.Status)
	}
	f.s.jobs.schedule(f.s.hub)
	f.mu.Lock()
	sent := append([]Msg(nil), f.sent...)
	f.mu.Unlock()
	if len(sent) != 1 || sent[0].T != "job" || sent[0].ID != j.ID || sent[0].Method != "image.generate" || !strings.Contains(string(sent[0].Params), "a cat") {
		t.Fatalf("the machine was not given the job: %+v", sent)
	}
	if got, _ := f.s.jobs.get(j.ID); got.Status != "running" || got.Machine != "pgx" {
		t.Fatalf("job = %+v", got)
	}
	// the machine is busy: a second job waits
	j2 := decodeJob(t, f.submit("image.generate", `{"prompt":"dog"}`, nil))
	f.s.jobs.schedule(f.s.hub)
	if got, _ := f.s.jobs.get(j2.ID); got.Status != "queued" || got.Why != "busy" {
		t.Fatalf("one job per machine: %+v", got)
	}
	// progress, then the result is uploaded with the machine's own token
	pb, _ := json.Marshal(map[string]any{"status": "running", "pct": 0.5})
	f.s.jobs.onAgentMsg("pgx", Msg{T: "job_progress", ID: j.ID, Data: pb})
	if got, _ := f.s.jobs.get(j.ID); got.Pct != 0.5 || got.Detail != "running" {
		t.Errorf("progress not recorded: %+v", got)
	}
	if w := f.agentReq("tok-other", "POST", "job="+j.ID+"&name=a.png&dir=out", strings.NewReader("PNG")); w.Code != 404 {
		t.Errorf("another machine must not write into the job: %d", w.Code)
	}
	if w := f.agentReq("bad", "POST", "job="+j.ID+"&name=a.png&dir=out", strings.NewReader("PNG")); w.Code != 401 {
		t.Errorf("a wrong token: %d", w.Code)
	}
	if w := f.agentReq("tok-pgx", "POST", "job="+j.ID+"&name=../../evil.png&dir=out", strings.NewReader("PNG")); w.Code == 200 {
		if _, err := os.Stat(filepath.Join(f.s.jobs.jobDir(j.ID), "out", "evil.png")); err != nil {
			t.Error("a path in the name was not reduced to a plain file name")
		}
	}
	if w := f.agentReq("tok-pgx", "POST", "job="+j.ID+"&name=a.png&dir=out", strings.NewReader("PNGDATA")); w.Code != 200 {
		t.Fatalf("upload: %d %s", w.Code, w.Body.String())
	}
	res, _ := json.Marshal(map[string]any{"outputs": []JobOutput{{Name: "a.png", Size: 7, Mime: "image/png"}, {Name: "ghost.png", Mime: "image/png"}}})
	f.s.jobs.onAgentMsg("pgx", Msg{T: "job_done", ID: j.ID, OK: true, Status: 200, Data: res})
	got, _ := f.s.jobs.get(j.ID)
	if got.Status != "done" || len(got.Outputs) != 1 || got.Outputs[0].Size != 7 {
		t.Fatalf("only files that really arrived count: %+v", got)
	}
	out := f.do("GET", "/studio/api/jobs/"+j.ID+"/out/a.png", nil, "", true)
	if out.Code != 200 || out.Body.String() != "PNGDATA" || out.Header().Get("Content-Type") != "image/png" {
		t.Errorf("download: %d %q %s", out.Code, out.Body.String(), out.Header().Get("Content-Type"))
	}
	if f.do("GET", "/studio/api/jobs/"+j.ID+"/out/a.png", nil, "", false).Code != 401 {
		t.Error("results are private to a logged-in user")
	}
	// the machine is free again: the waiting job goes out
	f.s.jobs.schedule(f.s.hub)
	if got, _ := f.s.jobs.get(j2.ID); got.Status != "running" {
		t.Errorf("the next job did not start: %+v", got)
	}
}

func TestInputFilesTravelToTheMachine(t *testing.T) {
	f := newStudioFixture(t)
	f.machine("pgx", "available", "image.edit")
	w := f.submit("image.edit", `{"prompt":"make it blue"}`, map[string][2]string{"image": {"../../my pic.png", "PICDATA"}})
	if w.Code != 200 {
		t.Fatalf("%d %s", w.Code, w.Body.String())
	}
	j := decodeJob(t, w)
	name, _ := j.Input["image"].(string)
	if name == "" || strings.ContainsAny(name, "/\\") {
		t.Fatalf("the stored name must be plain: %q", name)
	}
	f.s.jobs.schedule(f.s.hub)
	if r := f.agentReq("tok-pgx", "GET", "job="+j.ID+"&name="+url.QueryEscape(name)+"&dir=in", nil); r.Code != 200 || r.Body.String() != "PICDATA" {
		t.Errorf("the machine could not fetch the input: %d %q", r.Code, r.Body.String())
	}
	if r := f.agentReq("tok-other", "GET", "job="+j.ID+"&name="+url.QueryEscape(name)+"&dir=in", nil); r.Code != 404 {
		t.Errorf("another machine got the input: %d", r.Code)
	}
	// a missing required file, an unexpected file, a file that is too large
	if w := f.submit("image.edit", `{"prompt":"x"}`, nil); w.Code != 400 {
		t.Errorf("no picture: %d", w.Code)
	}
	if w := f.submit("image.generate", `{"prompt":"x"}`, map[string][2]string{"image": {"a.png", "x"}}); w.Code != 400 {
		t.Errorf("a file for a capability that takes none: %d", w.Code)
	}
	big := strings.Repeat("x", 20_000_001)
	f.machine("big", "available", "image.edit")
	if w := f.submit("image.edit", `{"prompt":"x"}`, map[string][2]string{"image": {"a.png", big}}); w.Code != 413 {
		t.Errorf("too large: %d", w.Code)
	}
}

func TestOwnerModeQueuesAndRetries(t *testing.T) {
	f := newStudioFixture(t)
	f.machine("pgx", "owner", "image.generate")
	j := decodeJob(t, f.submit("image.generate", `{"prompt":"x"}`, nil))
	f.s.jobs.schedule(f.s.hub)
	if got, _ := f.s.jobs.get(j.ID); got.Status != "queued" || got.Why != "owner" {
		t.Fatalf("while the owner is using the machine the job waits: %+v", got)
	}
	f.machine("pgx", "available", "image.generate")
	f.s.jobs.schedule(f.s.hub)
	if got, _ := f.s.jobs.get(j.ID); got.Status != "running" {
		t.Fatalf("it must start when the machine is free: %+v", got)
	}
	// the machine says "not now" (the owner came back): back in the queue, not failed
	f.s.jobs.onAgentMsg("pgx", Msg{T: "job_done", ID: j.ID, Status: 409, Err: "stopped because the owner is using the machine"})
	if got, _ := f.s.jobs.get(j.ID); got.Status != "queued" {
		t.Fatalf("a bounced job must be queued again: %+v", got)
	}
	// a real failure is reported as such
	f.s.jobs.mu.Lock()
	f.s.jobs.jobs[j.ID].notBefore = time.Time{}
	f.s.jobs.mu.Unlock()
	f.s.jobs.schedule(f.s.hub)
	f.s.jobs.onAgentMsg("pgx", Msg{T: "job_done", ID: j.ID, Status: 500, Err: "the engine failed: out of memory"})
	if got, _ := f.s.jobs.get(j.ID); got.Status != "failed" || !strings.Contains(got.Err, "out of memory") {
		t.Fatalf("%+v", got)
	}
}

func TestADisconnectedMachineGivesTheJobBack(t *testing.T) {
	f := newStudioFixture(t)
	f.machine("pgx", "available", "image.generate")
	j := decodeJob(t, f.submit("image.generate", `{"prompt":"x"}`, nil))
	f.s.jobs.schedule(f.s.hub)
	f.s.jobs.machineGone("pgx")
	if got, _ := f.s.jobs.get(j.ID); got.Status != "queued" || got.Machine != "" {
		t.Fatalf("%+v", got)
	}
	if len(f.s.jobs.busy) != 0 {
		t.Error("the machine is still marked busy")
	}
}

func TestCancelAndDelete(t *testing.T) {
	f := newStudioFixture(t)
	f.machine("pgx", "available", "image.generate")
	a := decodeJob(t, f.submit("image.generate", `{"prompt":"a"}`, nil))
	b := decodeJob(t, f.submit("image.generate", `{"prompt":"b"}`, nil))
	f.s.jobs.schedule(f.s.hub) // a runs, b waits
	if w := f.do("POST", "/studio/api/jobs/"+b.ID+"/cancel", nil, "", true); w.Code != 200 {
		t.Fatal(w.Code)
	}
	if got, _ := f.s.jobs.get(b.ID); got.Status != "cancelled" {
		t.Errorf("a waiting job is cancelled at once: %+v", got)
	}
	f.do("POST", "/studio/api/jobs/"+a.ID+"/cancel", nil, "", true)
	f.mu.Lock()
	last := f.sent[len(f.sent)-1]
	f.mu.Unlock()
	if last.T != "job_cancel" || last.ID != a.ID {
		t.Errorf("the machine must be told to stop: %+v", last)
	}
	f.s.jobs.onAgentMsg("pgx", Msg{T: "job_done", ID: a.ID, Status: 499, Err: "cancelled"})
	if got, _ := f.s.jobs.get(a.ID); got.Status != "cancelled" {
		t.Errorf("%+v", got)
	}
	os.MkdirAll(filepath.Join(f.s.jobs.jobDir(b.ID), "out"), 0750)
	if w := f.do("DELETE", "/studio/api/jobs/"+b.ID, nil, "", true); w.Code != 200 {
		t.Fatal(w.Code)
	}
	if _, ok := f.s.jobs.get(b.ID); ok {
		t.Error("the job is still listed")
	}
	if _, err := os.Stat(f.s.jobs.jobDir(b.ID)); err == nil {
		t.Error("its files stayed")
	}
}

func TestQueueIsBoundedAndSurvivesARestart(t *testing.T) {
	f := newStudioFixture(t)
	f.machine("pgx", "owner", "image.generate") // nothing runs, the queue fills
	for i := 0; i < maxQueuedJobs; i++ {
		if w := f.submit("image.generate", `{"prompt":"x"}`, nil); w.Code != 200 {
			t.Fatalf("job %d: %d", i, w.Code)
		}
	}
	if w := f.submit("image.generate", `{"prompt":"x"}`, nil); w.Code != 429 {
		t.Errorf("the queue must be bounded: %d", w.Code)
	}
	again := newJobStore(f.s.jobs.dir, f.s.jobs.cat)
	if len(again.list("studio")) != maxQueuedJobs {
		t.Errorf("the queue was not kept: %d", len(again.list("studio")))
	}
}

func TestOldResultsExpire(t *testing.T) {
	f := newStudioFixture(t)
	f.machine("pgx", "owner", "image.generate")
	j := decodeJob(t, f.submit("image.generate", `{"prompt":"x"}`, nil))
	f.s.jobs.mu.Lock()
	jj := f.s.jobs.jobs[j.ID]
	jj.Status, jj.Ended = "done", time.Now().Add(-8*24*time.Hour).Unix()
	f.s.jobs.mu.Unlock()
	os.MkdirAll(filepath.Join(f.s.jobs.jobDir(j.ID), "out"), 0750)
	f.s.jobs.expire()
	if _, ok := f.s.jobs.get(j.ID); ok {
		t.Error("an 8 day old result must be gone")
	}
}

func TestOpeningATabWarmsTheEngineOnAMachineThatOffersIt(t *testing.T) {
	f := newStudioFixture(t)
	if w := f.do("POST", "/studio/api/warm", strings.NewReader(`{"capability":"image.generate"}`), "application/json", false); w.Code != 401 {
		t.Errorf("without a login: %d", w.Code)
	}
	if w := f.do("POST", "/studio/api/warm", strings.NewReader(`{"capability":"nope"}`), "application/json", true); w.Code != 400 {
		t.Errorf("an unknown capability: %d", w.Code)
	}
	if w := f.do("GET", "/studio/api/warm", nil, "", true); w.Code != 405 {
		t.Errorf("GET: %d", w.Code)
	}
	// nobody offers it
	w := f.do("POST", "/studio/api/warm", strings.NewReader(`{"capability":"image.generate"}`), "application/json", true)
	if !strings.Contains(w.Body.String(), `"warming":false`) {
		t.Errorf("no machine: %s", w.Body.String())
	}
	f.machine("pgx", "available", "image.generate")
	sentWarm := func() int {
		f.mu.Lock()
		defer f.mu.Unlock()
		n := 0
		for _, m := range f.sent {
			if m.T == "warm" && m.Method == "image.generate" {
				n++
			}
		}
		return n
	}
	w = f.do("POST", "/studio/api/warm", strings.NewReader(`{"capability":"image.generate"}`), "application/json", true)
	if !strings.Contains(w.Body.String(), `"warming":true`) || sentWarm() != 1 {
		t.Fatalf("the machine must be asked once: %s sent=%d", w.Body.String(), sentWarm())
	}
	// asked again at once: told it is warming, but the machine is not asked again
	f.do("POST", "/studio/api/warm", strings.NewReader(`{"capability":"image.generate"}`), "application/json", true)
	if sentWarm() != 1 {
		t.Errorf("the machine was asked again within the minute: %d", sentWarm())
	}
	// later, again
	f.s.jobs.mu.Lock()
	f.s.jobs.warmAt["pgx|image.generate"] = time.Now().Add(-time.Minute)
	f.s.jobs.mu.Unlock()
	f.do("POST", "/studio/api/warm", strings.NewReader(`{"capability":"image.generate"}`), "application/json", true)
	if sentWarm() != 2 {
		t.Errorf("after the pause it must ask again: %d", sentWarm())
	}
	// a machine the owner is using is left alone
	f.machine("pgx", "owner", "image.generate")
	f.s.jobs.mu.Lock()
	f.s.jobs.warmAt = nil
	f.s.jobs.mu.Unlock()
	w = f.do("POST", "/studio/api/warm", strings.NewReader(`{"capability":"image.generate"}`), "application/json", true)
	if !strings.Contains(w.Body.String(), `"warming":false`) || sentWarm() != 2 {
		t.Errorf("owner mode: %s sent=%d", w.Body.String(), sentWarm())
	}
}

func TestStudioTellsHowMuchMemoryAJobUses(t *testing.T) {
	f := newStudioFixture(t)
	f.machine("pgx", "available", "image.generate")
	var caps struct{ Capabilities []studioCap }
	json.Unmarshal(f.do("GET", "/studio/api/capabilities", nil, "", true).Body.Bytes(), &caps)
	for _, c := range caps.Capabilities {
		if c.MemBytes <= 0 {
			t.Errorf("%s: the page needs the memory figure for its footnote", c.ID)
		}
	}
}

func TestAWaitForTheOwnerSaysWhyAndForHowLong(t *testing.T) {
	f := newStudioFixture(t)
	f.machine("pgx", "available", "image.generate")
	// the machine reports that its owner is at the desktop and will be gone in two minutes if left alone
	b, _ := json.Marshal(map[string]any{"ts": time.Now().Unix(), "mode": "owner", "mode_kind": "desktop", "mode_retry_s": 120, "modules": []string{"image.generate"}, "mem_avail": 90e9})
	f.s.hub.recordStats("pgx", b)
	var caps struct{ Capabilities []studioCap }
	json.Unmarshal(f.do("GET", "/studio/api/capabilities", nil, "", true).Body.Bytes(), &caps)
	var c studioCap
	for _, x := range caps.Capabilities {
		if x.ID == "image.generate" {
			c = x
		}
	}
	if c.Available || c.Why != "owner" || c.WhyKind != "desktop" || c.RetryS < 100 || c.RetryS > 125 {
		t.Fatalf("the page needs the reason and the wait: %+v", c)
	}
	j := decodeJob(t, f.submit("image.generate", `{"prompt":"x"}`, nil))
	f.s.jobs.schedule(f.s.hub)
	got, _ := f.s.jobs.get(j.ID)
	if got.Status != "queued" || got.Why != "owner" || got.WhyKind != "desktop" || got.RetryS < 100 {
		t.Fatalf("the waiting job must carry it too: %+v", got)
	}
	// another program on the GPU
	b, _ = json.Marshal(map[string]any{"ts": time.Now().Unix(), "mode": "owner", "mode_kind": "gpu", "modules": []string{"image.generate"}, "mem_avail": 90e9})
	f.s.hub.recordStats("pgx", b)
	f.s.jobs.schedule(f.s.hub)
	if got, _ := f.s.jobs.get(j.ID); got.WhyKind != "gpu" || got.RetryS != 0 {
		t.Errorf("an unknown wait is reported as unknown: %+v", got)
	}
	// free again: the job starts and the reason is cleared
	f.machine("pgx", "available", "image.generate")
	f.s.jobs.schedule(f.s.hub)
	if got, _ := f.s.jobs.get(j.ID); got.Status != "running" || got.Why != "" || got.WhyKind != "" || got.RetryS != 0 {
		t.Errorf("%+v", got)
	}
}