package main

import (
	"encoding/base64"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/gorilla/websocket"
)

// comfyFixture is a server with one machine whose "agent" answers tunnel messages in-process.
type comfyFixture struct {
	f    *studioFixture
	srv  *httptest.Server
	mu   sync.Mutex
	reqs []Msg
	body strings.Builder
}

func newComfyFixture(t *testing.T) *comfyFixture {
	f := newStudioFixture(t)
	f.s.comfySess = newSessionStore("").set("comfy")
	f.s.comfySess.add("comfy-token", time.Hour)
	cf := &comfyFixture{f: f}
	c := &conn{sink: func(m Msg) {
		cf.mu.Lock()
		cf.reqs = append(cf.reqs, m)
		cf.mu.Unlock()
		switch m.T {
		case "tun_req":
			go func() {
				var p struct {
					Method  string
					Headers map[string]string
				}
				json.Unmarshal(m.Params, &p)
				ctype, content := "text/plain", "hello "+m.D+" "
				if m.D == "/page" {
					ctype, content = "text/html; charset=utf-8", `<html><head><link rel="icon" type="image/x-icon" href="/favicon.ico"><title>ComfyUI</title></head><body><div id="app"></div></body></html>`
				}
				hdr, _ := json.Marshal(map[string]string{"Content-Type": ctype, "X-Seen-Cookie": p.Headers["Cookie"], "X-Seen-Origin": p.Headers["Origin"]})
				f.s.tuns.deliver("pgx", Msg{T: "tun_res", ID: m.ID, Status: 200, Params: hdr})
				f.s.tuns.deliver("pgx", Msg{T: "tun_data", ID: m.ID, D: base64.StdEncoding.EncodeToString([]byte(content))})
				f.s.tuns.deliver("pgx", Msg{T: "tun_data", ID: m.ID, D: base64.StdEncoding.EncodeToString([]byte(p.Method))})
				f.s.tuns.deliver("pgx", Msg{T: "tun_end", ID: m.ID})
			}()
		case "tun_body":
			if raw, err := base64.StdEncoding.DecodeString(m.D); err == nil {
				cf.mu.Lock()
				cf.body.Write(raw)
				cf.mu.Unlock()
			}
		case "tun_ws_open":
			f.s.tuns.deliver("pgx", Msg{T: "tun_ws_msg", ID: m.ID, Status: websocket.TextMessage, D: base64.StdEncoding.EncodeToString([]byte("welcome"))})
		case "tun_ws_msg": // echo
			f.s.tuns.deliver("pgx", Msg{T: "tun_ws_msg", ID: m.ID, Status: m.Status, D: m.D})
		}
	}}
	f.s.hub.mu.Lock()
	f.s.hub.agents["pgx"] = c
	f.s.hub.mu.Unlock()
	b, _ := json.Marshal(map[string]any{"ts": time.Now().Unix(), "mode": "available", "modules": []string{"image.generate"}, "mem_avail": 90e9})
	f.s.hub.recordStats("pgx", b)
	cf.srv = httptest.NewServer(f.s.comfyMux())
	t.Cleanup(cf.srv.Close)
	return cf
}

func (cf *comfyFixture) get(path string, cookie bool, accept string) (*http.Response, string) {
	req, _ := http.NewRequest("GET", cf.srv.URL+path, nil)
	if cookie {
		req.AddCookie(&http.Cookie{Name: "wxch_comfy", Value: "comfy-token"})
		req.AddCookie(&http.Cookie{Name: "other", Value: "secret"})
	}
	if accept != "" {
		req.Header.Set("Accept", accept)
	}
	req.Header.Set("Origin", "https://comfy.wxch.org")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		panic(err)
	}
	defer resp.Body.Close()
	b, _ := io.ReadAll(resp.Body)
	return resp, string(b)
}

func TestComfyPageNeedsLogin(t *testing.T) {
	cf := newComfyFixture(t)
	if r, body := cf.get("/", false, "text/html"); r.StatusCode != 200 || !strings.Contains(body, "comfy-login") {
		t.Errorf("a browser without a session must get the login page: %d", r.StatusCode)
	}
	if r, _ := cf.get("/api/system_stats", false, ""); r.StatusCode != 401 {
		t.Errorf("API without a session: %d", r.StatusCode)
	}
	cf.mu.Lock()
	n := len(cf.reqs)
	cf.mu.Unlock()
	if n != 0 {
		t.Error("an unauthenticated request reached the machine")
	}
}

func TestComfyRequestsReachTheMachineWithoutOurCookies(t *testing.T) {
	cf := newComfyFixture(t)
	r, body := cf.get("/object_info?x=1", true, "")
	if r.StatusCode != 200 || body != "hello /object_info?x=1 GET" {
		t.Fatalf("%d %q", r.StatusCode, body)
	}
	if r.Header.Get("X-Seen-Cookie") != "" {
		t.Errorf("our session cookie was forwarded to the engine: %q", r.Header.Get("X-Seen-Cookie"))
	}
	// a POST body travels in chunks and ends with a marker
	req, _ := http.NewRequest("POST", cf.srv.URL+"/prompt", strings.NewReader(`{"prompt":{}}`))
	req.AddCookie(&http.Cookie{Name: "wxch_comfy", Value: "comfy-token"})
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	b, _ := io.ReadAll(resp.Body)
	resp.Body.Close()
	if string(b) != "hello /prompt POST" {
		t.Errorf("post answer %q", b)
	}
	cf.mu.Lock()
	defer cf.mu.Unlock()
	if cf.body.String() != `{"prompt":{}}` {
		t.Errorf("the body that reached the machine: %q", cf.body.String())
	}
	last := cf.reqs[len(cf.reqs)-1]
	if last.T != "tun_body" || last.Status != 1 {
		t.Errorf("the body must end with a marker: %+v", last)
	}
}

func TestComfyWebsocketIsCarried(t *testing.T) {
	cf := newComfyFixture(t)
	h := http.Header{"Cookie": {"wxch_comfy=comfy-token"}, "Origin": {cf.srv.URL}}
	ws, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(cf.srv.URL, "http")+"/ws?clientId=a", h)
	if err != nil {
		t.Fatal(err)
	}
	defer ws.Close()
	ws.SetReadDeadline(time.Now().Add(5 * time.Second))
	if _, m, err := ws.ReadMessage(); err != nil || string(m) != "welcome" {
		t.Fatalf("first frame: %q %v", m, err)
	}
	ws.WriteMessage(websocket.BinaryMessage, []byte{1, 2, 3})
	mt, m, err := ws.ReadMessage()
	if err != nil || mt != websocket.BinaryMessage || len(m) != 3 {
		t.Fatalf("echo: %d %v %v", mt, m, err)
	}
	if _, _, err := websocket.DefaultDialer.Dial("ws"+strings.TrimPrefix(cf.srv.URL, "http")+"/ws", http.Header{"Origin": {cf.srv.URL}}); err == nil {
		t.Error("a websocket without a session was accepted")
	}
}

func TestComfyNeedsAMachineWithAComfyModule(t *testing.T) {
	cf := newComfyFixture(t)
	b, _ := json.Marshal(map[string]any{"ts": time.Now().Unix(), "mode": "available", "modules": []string{"audio.transcribe"}})
	cf.f.s.hub.recordStats("pgx", b) // only Whisper installed: no ComfyUI
	if r, _ := cf.get("/", true, ""); r.StatusCode != 503 {
		t.Errorf("a machine without a ComfyUI module must not be offered: %d", r.StatusCode)
	}
}

func TestComfyTunnelEndsWhenTheMachineDisconnects(t *testing.T) {
	cf := newComfyFixture(t)
	cf.f.s.hub.mu.Lock()
	cf.f.s.hub.agents["pgx"] = &conn{sink: func(m Msg) {
		if m.T == "tun_req" {
			go func() {
				time.Sleep(100 * time.Millisecond)
				cf.f.s.tuns.agentGone("pgx")
			}()
		}
	}}
	cf.f.s.hub.mu.Unlock()
	r, body := cf.get("/x", true, "")
	if r.StatusCode != 502 || !strings.Contains(body, "disconnected") {
		t.Errorf("%d %q", r.StatusCode, body)
	}
}

func TestComfyPageGetsOurLogoAndKeepsItsOwnContent(t *testing.T) {
	cf := newComfyFixture(t)
	req, _ := http.NewRequest("GET", cf.srv.URL+"/page", nil)
	req.Host = "comfy.wxch.org"
	req.AddCookie(&http.Cookie{Name: "wxch_comfy", Value: "comfy-token"})
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		t.Fatal(err)
	}
	b, _ := io.ReadAll(resp.Body)
	resp.Body.Close()
	page := string(b)
	if resp.StatusCode != 200 || !strings.Contains(page, "<title>ComfyUI</title>") || !strings.Contains(page, `<div id="app">`) {
		t.Fatalf("the page itself must survive: %d %q", resp.StatusCode, page)
	}
	if strings.Contains(page, "favicon.ico") || strings.Count(page, `rel="icon"`) != 1 || !strings.Contains(page, "image/svg+xml") {
		t.Errorf("our icon must replace the page's own: %q", page)
	}
	if !strings.Contains(page, `href="https://studio.wxch.org/"`) || !strings.Contains(page, "<polygon") {
		t.Errorf("the corner mark linking back to the studio is missing: %q", page)
	}
	if !strings.Contains(page, "/_wxch/fetch") || strings.Count(page, "<script>(function(){") != 1 {
		t.Errorf("the script that sends model downloads to the machine is missing: %q", page)
	}
	if resp.Header.Get("Content-Length") == "" {
		t.Error("the changed page needs its new length")
	}
	// anything that is not HTML passes through untouched
	if r, body := cf.get("/object_info", true, ""); strings.Contains(body, "<svg") || r.StatusCode != 200 {
		t.Errorf("non-HTML answers must not be changed: %d %q", r.StatusCode, body)
	}
}

func TestComfyMessagePagesCarryTheLogo(t *testing.T) {
	cf := newComfyFixture(t)
	// no machine with a ComfyUI module
	b, _ := json.Marshal(map[string]any{"ts": time.Now().Unix(), "mode": "available", "modules": []string{"audio.transcribe"}})
	cf.f.s.hub.recordStats("pgx", b)
	r, body := cf.get("/", true, "text/html")
	if r.StatusCode != 503 || !strings.Contains(body, `rel="icon"`) || !strings.Contains(body, "<polygon") {
		t.Errorf("the no-machine page: %d %q", r.StatusCode, body)
	}
	// the page the agent shows while the engine starts keeps its 503 and gets the logo
	b, _ = json.Marshal(map[string]any{"ts": time.Now().Unix(), "mode": "available", "modules": []string{"image.generate"}})
	cf.f.s.hub.recordStats("pgx", b)
	cf.f.s.hub.mu.Lock()
	cf.f.s.hub.agents["pgx"] = &conn{sink: func(m Msg) {
		if m.T == "tun_req" {
			go func() {
				h, _ := json.Marshal(map[string]string{"Content-Type": "text/html; charset=utf-8"})
				cf.f.s.tuns.deliver("pgx", Msg{T: "tun_res", ID: m.ID, Status: 503, Params: h})
				cf.f.s.tuns.deliver("pgx", Msg{T: "tun_data", ID: m.ID, D: base64.StdEncoding.EncodeToString([]byte(`<!doctype html><meta http-equiv="refresh" content="5"><title>starting</title><body>starting</body>`))})
				cf.f.s.tuns.deliver("pgx", Msg{T: "tun_end", ID: m.ID})
			}()
		}
	}}
	cf.f.s.hub.mu.Unlock()
	r, body = cf.get("/", true, "text/html")
	if r.StatusCode != 503 || !strings.Contains(body, `refresh`) || !strings.Contains(body, `rel="icon"`) || !strings.Contains(body, "<polygon") {
		t.Errorf("the starting page: %d %q", r.StatusCode, body)
	}
	// an error from the tunnel is a page of ours too
	cf.f.s.hub.mu.Lock()
	cf.f.s.hub.agents["pgx"] = &conn{sink: func(m Msg) {
		if m.T == "tun_req" {
			go func() { cf.f.s.tuns.deliver("pgx", Msg{T: "tun_end", ID: m.ID, Err: "engine exploded"}) }()
		}
	}}
	cf.f.s.hub.mu.Unlock()
	r, body = cf.get("/", true, "text/html")
	if r.StatusCode != 502 || !strings.Contains(body, "engine exploded") || !strings.Contains(body, "<polygon") {
		t.Errorf("the error page: %d %q", r.StatusCode, body)
	}
}