package main

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

func TestCleanFetchURL(t *testing.T) {
	good := map[string]string{
		"https://huggingface.co/Comfy-Org/x/resolve/main/split_files/vae/ae.safetensors": "https://huggingface.co/Comfy-Org/x/resolve/main/split_files/vae/ae.safetensors",
		"https://huggingface.co/Comfy-Org/x/blob/main/split_files/vae/ae.safetensors":    "https://huggingface.co/Comfy-Org/x/resolve/main/split_files/vae/ae.safetensors",
		"https://huggingface.co/a/b/resolve/main/m.safetensors?download=true#frag":       "https://huggingface.co/a/b/resolve/main/m.safetensors",
		"https://cdn-lfs.huggingface.co/a/b/resolve/main/m.safetensors":                  "https://cdn-lfs.huggingface.co/a/b/resolve/main/m.safetensors",
	}
	for in, want := range good {
		got, _, err := cleanFetchURL(in)
		if err != nil || got != want {
			t.Errorf("%s -> %q, %v (want %s)", in, got, err, want)
		}
	}
	for _, in := range []string{
		"http://huggingface.co/a/b/resolve/main/m.safetensors",
		"https://evilhuggingface.co/a/b/resolve/main/m.safetensors",
		"https://huggingface.co.evil.org/a/b/resolve/main/m.safetensors",
		"https://user:pw@huggingface.co/a/b/resolve/main/m.safetensors",
		"https://huggingface.co:8443/a/b/resolve/main/m.safetensors",
		"https://huggingface.co/a/b",
		"https://example.org/a/resolve/main/m.safetensors",
		"file:///etc/passwd", "", "not a url",
	} {
		if got, _, err := cleanFetchURL(in); err == nil {
			t.Errorf("%q must be refused, got %s", in, got)
		}
	}
}

func fetchSetup(t *testing.T) (nodeConfig, string) {
	cfg, _, cat := moduleSetup(t)
	e, _ := cat.Engine("comfyui")
	writeModuleState(cfg, moduleState{Installed: map[string]string{"a.one": "x"}})
	root := filepath.Join(engineDir(cfg, e), "ComfyUI", "models")
	os.MkdirAll(root, 0755)
	fetches.Lock()
	fetches.jobs = map[string]*fetchJob{}
	fetches.Unlock()
	oldSize, oldDl := fetchSizeFn, fetchDownload
	t.Cleanup(func() { fetchSizeFn, fetchDownload = oldSize, oldDl })
	return cfg, root
}

func waitFetch(t *testing.T, key string) fetchJob {
	t.Helper()
	for i := 0; i < 200; i++ {
		fetches.Lock()
		j := fetches.jobs[key]
		var c fetchJob
		if j != nil {
			c = *j
		}
		fetches.Unlock()
		if c.State == "done" || c.State == "failed" {
			return c
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatal("the download did not finish")
	return fetchJob{}
}

func TestFetchDownloadsIntoTheModelsFolder(t *testing.T) {
	cfg, root := fetchSetup(t)
	fetchSizeFn = func(context.Context, string) (int64, error) { return 1234, nil }
	var gotURL, gotDest string
	fetchDownload = func(_ context.Context, u, dest string, size int64, on func(int64)) error {
		gotURL, gotDest = u, dest
		on(600)
		on(634)
		return os.WriteFile(dest, make([]byte, size), 0644)
	}
	os.MkdirAll(filepath.Join(root, "vae"), 0755)
	j, err := startFetch(cfg, fetchReq{URL: "https://huggingface.co/a/b/blob/main/split_files/vae/ae.safetensors?download=true", Dir: "vae"})
	if err != nil || j.Name != "ae.safetensors" {
		t.Fatalf("%+v %v", j, err)
	}
	done := waitFetch(t, "vae/ae.safetensors")
	if done.State != "done" || done.Done != 1234 || done.Total != 1234 {
		t.Fatalf("%+v", done)
	}
	if gotURL != "https://huggingface.co/a/b/resolve/main/split_files/vae/ae.safetensors" || gotDest != filepath.Join(root, "vae", "ae.safetensors") {
		t.Errorf("downloaded %s into %s", gotURL, gotDest)
	}
	// asking again for a file that is there does nothing
	calls := 0
	fetchDownload = func(context.Context, string, string, int64, func(int64)) error { calls++; return nil }
	startFetch(cfg, fetchReq{URL: "https://huggingface.co/a/b/resolve/main/vae/ae.safetensors", Dir: "vae"})
	waitFetch(t, "vae/ae.safetensors")
	time.Sleep(50 * time.Millisecond)
	if calls != 0 {
		t.Error("a file that is already there was downloaded again")
	}
}

func TestFetchRefusesWhatItShould(t *testing.T) {
	cfg, _ := fetchSetup(t)
	ok := "https://huggingface.co/a/b/resolve/main/x.safetensors"
	for name, req := range map[string]fetchReq{
		"a host that is not Hugging Face": {URL: "https://example.org/a/resolve/main/x.safetensors", Dir: "vae"},
		"a folder that climbs out":        {URL: ok, Dir: "../etc"},
		"a folder with a slash":           {URL: ok, Dir: "vae/../../x"},
		"no folder":                       {URL: ok},
		"a name with a path":              {URL: ok, Dir: "vae", Name: "../../evil.safetensors"},
		"a file that is not a model":      {URL: "https://huggingface.co/a/b/resolve/main/run.sh", Dir: "vae"},
		"an upper case folder":            {URL: ok, Dir: "VAE"},
	} {
		if _, err := startFetch(cfg, req); err == nil {
			t.Errorf("%s was accepted", name)
		}
	}
	// no module installed: nothing to download for
	writeModuleState(cfg, moduleState{Installed: map[string]string{}})
	if _, err := startFetch(cfg, fetchReq{URL: ok, Dir: "vae"}); err == nil || !strings.Contains(err.Error(), "no image or video module") {
		t.Errorf("got %v", err)
	}
}

func TestFetchReportsFailuresAndTooLittleDisk(t *testing.T) {
	cfg, _ := fetchSetup(t)
	fetchSizeFn = func(context.Context, string) (int64, error) {
		return 0, fmt.Errorf("the file was not found on Hugging Face")
	}
	startFetch(cfg, fetchReq{URL: "https://huggingface.co/a/b/resolve/main/x.safetensors", Dir: "vae"})
	if j := waitFetch(t, "vae/x.safetensors"); j.State != "failed" || !strings.Contains(j.Err, "not found") {
		t.Errorf("%+v", j)
	}
	if _, _, err := diskFree(cfg.ModulesDir); err != nil {
		return // the disk size can only be read on Linux
	}
	cfg.MinFreeGB = 1 << 20 // far more than any disk has
	fetchSizeFn = func(context.Context, string) (int64, error) { return 1e9, nil }
	startFetch(cfg, fetchReq{URL: "https://huggingface.co/a/b/resolve/main/y.safetensors", Dir: "vae"})
	if j := waitFetch(t, "vae/y.safetensors"); j.State != "failed" || !strings.Contains(j.Err, "not enough free disk") {
		t.Errorf("%+v", j)
	}
}

func TestTheTunnelAnswersFetchRequestsItself(t *testing.T) {
	cfg, root := fetchSetup(t)
	os.MkdirAll(filepath.Join(root, "vae"), 0755)
	cfgOverride = &cfg
	defer func() { cfgOverride = nil }()
	fetchSizeFn = func(context.Context, string) (int64, error) { return 10, nil }
	fetchDownload = func(_ context.Context, _, dest string, size int64, on func(int64)) error {
		on(size)
		return os.WriteFile(dest, make([]byte, size), 0644)
	}
	run := func(method, path string, body string) (int, map[string]any) {
		var out []Msg
		in := make(chan Msg, 4)
		if body != "" {
			in <- Msg{T: "tun_body", D: base64.StdEncoding.EncodeToString([]byte(body)), Status: 1}
		}
		serveLocalTun(context.Background(), Msg{ID: "1", D: path}, &agentTun{in: in}, method, func(m Msg) { out = append(out, m) })
		var code int
		var data []byte
		for _, m := range out {
			switch m.T {
			case "tun_res":
				code = m.Status
			case "tun_data":
				b, _ := base64.StdEncoding.DecodeString(m.D)
				data = append(data, b...)
			}
		}
		var v map[string]any
		json.Unmarshal(data, &v)
		return code, v
	}
	code, v := run("POST", "/_wxch/fetch", `{"url":"https://huggingface.co/a/b/resolve/main/vae/q.safetensors","dir":"vae"}`)
	if code != 200 || v["ok"] != true {
		t.Fatalf("%d %v", code, v)
	}
	waitFetch(t, "vae/q.safetensors")
	code, v = run("GET", "/_wxch/fetch", "")
	if jobs, _ := v["jobs"].([]any); code != 200 || len(jobs) != 1 {
		t.Errorf("list: %d %v", code, v)
	}
	if code, _ = run("POST", "/_wxch/fetch", `{"url":"https://example.org/x.safetensors","dir":"vae"}`); code != 400 {
		t.Errorf("a bad address gave %d", code)
	}
	if code, _ = run("POST", "/_wxch/fetch", `not json`); code != 400 {
		t.Errorf("bad json gave %d", code)
	}
	if code, _ = run("GET", "/_wxch/other", ""); code != 404 {
		t.Errorf("unknown path gave %d", code)
	}
	if code, _ = run("DELETE", "/_wxch/fetch", ""); code != 405 {
		t.Errorf("DELETE gave %d", code)
	}
}