package main

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

	"github.com/gorilla/websocket"

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

type tunCollector struct {
	mu   sync.Mutex
	msgs []Msg
}

func (c *tunCollector) send(m Msg) { c.mu.Lock(); c.msgs = append(c.msgs, m); c.mu.Unlock() }
func (c *tunCollector) wait(t *testing.T, id, typ string) Msg {
	for i := 0; i < 200; i++ {
		c.mu.Lock()
		for _, m := range c.msgs {
			if m.ID == id && m.T == typ {
				c.mu.Unlock()
				return m
			}
		}
		c.mu.Unlock()
		time.Sleep(25 * time.Millisecond)
	}
	t.Fatalf("no %s for %s", typ, id)
	return Msg{}
}
func (c *tunCollector) data(id string) string {
	c.mu.Lock()
	defer c.mu.Unlock()
	var b strings.Builder
	for _, m := range c.msgs {
		if m.ID == id && m.T == "tun_data" {
			raw, _ := base64.StdEncoding.DecodeString(m.D)
			b.Write(raw)
		}
	}
	return b.String()
}

func tunSetup(t *testing.T, healthy bool) (*httptest.Server, *tunCollector) {
	for i := 0; i < 300; i++ { // an engine start from an earlier test may still be waiting for its health check
		if _, busy := startingEngine.Load("comfyui"); !busy {
			break
		}
		time.Sleep(20 * time.Millisecond)
	}
	cfg := nodeConfig{ModulesDir: t.TempDir()}
	writeModuleState(cfg, moduleState{Installed: map[string]string{"image.generate": "x"}})
	cfgOverride = &cfg
	eng := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/system_stats":
			w.Write([]byte("{}"))
		case "/ws":
			up := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return r.Header.Get("Origin") == "http://"+r.Host }}
			c, err := up.Upgrade(w, r, nil)
			if err != nil {
				return
			}
			defer c.Close()
			for {
				mt, m, err := c.ReadMessage()
				if err != nil {
					return
				}
				c.WriteMessage(mt, m)
			}
		default:
			b, _ := io.ReadAll(r.Body)
			w.Header().Set("X-Host", r.Host)
			w.Header().Set("X-Origin", r.Header.Get("Origin"))
			w.Header().Set("X-Cookie", r.Header.Get("Cookie"))
			w.Header().Set("Content-Type", "text/plain")
			w.Write([]byte(r.Method + " " + r.URL.RequestURI() + " body=" + string(b)))
		}
	}))
	base := eng.URL
	if !healthy {
		base = "http://127.0.0.1:1"
	}
	old := engineBase
	engineBase = func(*modules.Catalog, modules.Engine) string { return base }
	oldWait := engineWait
	engineWait = time.Second
	modRunnerOverride = &fakeModRunner{image: map[string]bool{"comfyui": true}}
	lastEngineStart.Delete("comfyui")
	t.Cleanup(func() {
		eng.Close()
		engineBase = old
		engineWait = oldWait
		cfgOverride = nil
		modRunnerOverride = nil
	})
	return eng, &tunCollector{}
}

func TestTunnelCarriesARequestToTheEngine(t *testing.T) {
	eng, col := tunSetup(t, true)
	params, _ := json.Marshal(map[string]any{"method": "POST", "headers": map[string]string{
		"Origin": "https://comfy.wxch.org", "Cookie": "wxch_comfy=secret", "Content-Type": "application/json", "Content-Length": "7"}})
	handleTunMsg(Msg{T: "tun_req", ID: "a", Method: "comfyui", D: "/prompt?x=1", Params: params}, col.send)
	handleTunMsg(Msg{T: "tun_body", ID: "a", D: base64.StdEncoding.EncodeToString([]byte("{\"a\":1}"))}, col.send)
	handleTunMsg(Msg{T: "tun_body", ID: "a", Status: 1}, col.send)
	res := col.wait(t, "a", "tun_res")
	col.wait(t, "a", "tun_end")
	if res.Status != 200 {
		t.Fatalf("status %d", res.Status)
	}
	if got := col.data("a"); got != `POST /prompt?x=1 body={"a":1}` {
		t.Errorf("answer %q", got)
	}
	var h map[string]string
	json.Unmarshal(res.Params, &h)
	host := strings.TrimPrefix(eng.URL, "http://")
	if h["X-Host"] != host || h["X-Origin"] != "http://"+host {
		t.Errorf("the engine must see its own address as Host and Origin, saw %q / %q", h["X-Host"], h["X-Origin"])
	}
	if h["X-Cookie"] != "" {
		t.Errorf("the session cookie reached the engine: %q", h["X-Cookie"])
	}
	if tunnelsOpen() {
		t.Error("the tunnel was not closed")
	}
}

func TestTunnelStartsAnEngineThatIsNotRunning(t *testing.T) {
	_, col := tunSetup(t, false)
	handleTunMsg(Msg{T: "tun_req", ID: "b", Method: "comfyui", D: "/", Params: []byte(`{"method":"GET"}`)}, col.send)
	if res := col.wait(t, "b", "tun_res"); res.Status != 503 {
		t.Fatalf("status %d", res.Status)
	}
	col.wait(t, "b", "tun_end")
	if !strings.Contains(col.data("b"), "refresh") {
		t.Errorf("the page must reload itself while the engine starts: %q", col.data("b"))
	}
	time.Sleep(200 * time.Millisecond)
	rn := modRunnerOverride.(*fakeModRunner)
	rn.mu.Lock()
	started := rn.running["comfyui"]
	rn.mu.Unlock()
	if !started {
		t.Error("the engine was not started")
	}
}

func TestTunnelRefusesAnEngineWithoutAnInstalledModule(t *testing.T) {
	_, col := tunSetup(t, true)
	writeModuleState(*cfgOverride, moduleState{Installed: map[string]string{"audio.transcribe": "x"}})
	handleTunMsg(Msg{T: "tun_req", ID: "c", Method: "comfyui", D: "/", Params: []byte(`{"method":"GET"}`)}, col.send)
	if res := col.wait(t, "c", "tun_res"); res.Status != 502 {
		t.Fatalf("status %d", res.Status)
	}
	handleTunMsg(Msg{T: "tun_req", ID: "d", Method: "whisper", D: "/", Params: []byte(`{"method":"GET"}`)}, col.send)
	if res := col.wait(t, "d", "tun_res"); res.Status != 200 && res.Status != 503 && res.Status != 502 {
		t.Fatalf("status %d", res.Status)
	}
}

func TestTunnelCarriesAWebsocket(t *testing.T) {
	_, col := tunSetup(t, true)
	handleTunMsg(Msg{T: "tun_ws_open", ID: "w", Method: "comfyui", D: "/ws?clientId=z"}, col.send)
	handleTunMsg(Msg{T: "tun_ws_msg", ID: "w", Status: websocket.TextMessage, D: base64.StdEncoding.EncodeToString([]byte("ping"))}, col.send)
	m := col.wait(t, "w", "tun_ws_msg")
	if raw, _ := base64.StdEncoding.DecodeString(m.D); string(raw) != "ping" || m.Status != websocket.TextMessage {
		t.Fatalf("echo %q %d", raw, m.Status)
	}
	if !tunnelsOpen() {
		t.Error("an open websocket must keep the engine from being stopped as idle")
	}
	handleTunMsg(Msg{T: "tun_ws_close", ID: "w"}, col.send)
	for i := 0; i < 100 && tunnelsOpen(); i++ {
		time.Sleep(20 * time.Millisecond)
	}
	if tunnelsOpen() {
		t.Error("the tunnel stayed open after the close")
	}
}

func TestStartingPageShowsTheEngineLogAndDoesNotRestartItInALoop(t *testing.T) {
	_, col := tunSetup(t, false)
	rn := modRunnerOverride.(*fakeModRunner)
	rn.logs = "Collecting torchsde" + string(rune(10)) + "ERROR: no matching distribution for <comfy-kitchen>"
	rn.exited = map[string]bool{"comfyui": true}
	// the first look starts the engine (it is not running)
	handleTunMsg(Msg{T: "tun_req", ID: "s1", Method: "comfyui", D: "/", Params: []byte(`{"method":"GET"}`)}, col.send)
	col.wait(t, "s1", "tun_end")
	page := col.data("s1")
	if !strings.Contains(page, "exited") || !strings.Contains(page, "no matching distribution") || strings.Contains(page, "<comfy-kitchen>") || !strings.Contains(page, "&lt;comfy-kitchen&gt;") {
		t.Errorf("the page must say what the engine is doing, with its log escaped: %q", page)
	}
	time.Sleep(150 * time.Millisecond)
	rn.mu.Lock()
	starts := 0
	if rn.running["comfyui"] {
		starts++
	}
	rn.mu.Unlock()
	if starts != 1 {
		t.Fatalf("the engine should have been started once, running=%d", starts)
	}
	// it dies again at once; reloading the page within a minute must not start it again
	rn.mu.Lock()
	delete(rn.running, "comfyui")
	rn.mu.Unlock()
	handleTunMsg(Msg{T: "tun_req", ID: "s2", Method: "comfyui", D: "/", Params: []byte(`{"method":"GET"}`)}, col.send)
	col.wait(t, "s2", "tun_end")
	time.Sleep(150 * time.Millisecond)
	rn.mu.Lock()
	again := rn.running["comfyui"]
	rn.mu.Unlock()
	if again {
		t.Error("an engine that just died was started again on the next page load")
	}
	// an engine that is running but not answering yet is left alone
	rn.mu.Lock()
	rn.exited = nil
	rn.running["comfyui"] = true
	rn.mu.Unlock()
	lastEngineStart.Delete("comfyui")
	handleTunMsg(Msg{T: "tun_req", ID: "s3", Method: "comfyui", D: "/", Params: []byte(`{"method":"GET"}`)}, col.send)
	col.wait(t, "s3", "tun_end")
	if _, ok := lastEngineStart.Load("comfyui"); ok {
		t.Error("an engine that is starting up was started again")
	}
	if !strings.Contains(col.data("s3"), "running") {
		t.Errorf("%q", col.data("s3"))
	}
}

func TestStartingPageSaysWhyTheEngineCouldNotBeStarted(t *testing.T) {
	_, col := tunSetup(t, false)
	lastEngineErr.Delete("comfyui")
	rn := modRunnerOverride.(*fakeModRunner)
	rn.startErr = "docker: Error response from daemon: could not select device driver <nvidia>"
	handleTunMsg(Msg{T: "tun_req", ID: "e1", Method: "comfyui", D: "/", Params: []byte(`{"method":"GET"}`)}, col.send)
	col.wait(t, "e1", "tun_end")
	for i := 0; i < 100; i++ { // the start runs in the background
		if _, ok := lastEngineErr.Load("comfyui"); ok {
			break
		}
		time.Sleep(20 * time.Millisecond)
	}
	lastEngineStart.Delete("comfyui") // a minute later
	handleTunMsg(Msg{T: "tun_req", ID: "e2", Method: "comfyui", D: "/", Params: []byte(`{"method":"GET"}`)}, col.send)
	col.wait(t, "e2", "tun_end")
	page := col.data("e2")
	if !strings.Contains(page, "could not be started") || !strings.Contains(page, "could not select device driver") || strings.Contains(page, "<nvidia>") {
		t.Errorf("the page must say why, escaped: %q", page)
	}
	lastEngineErr.Delete("comfyui")
}