package main

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync/atomic"
	"testing"
	"time"
)

// fakeHF serves a tiny "repo" with the same endpoints the downloader uses.
func fakeHF(t *testing.T, files map[string]string, rangeCount *int32) *httptest.Server {
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch {
		case r.URL.Path == "/api/models/org/model":
			type sib struct {
				Name string `json:"rfilename"`
				Size int64  `json:"size"`
			}
			var sibs []sib
			for n, c := range files {
				sibs = append(sibs, sib{n, int64(len(c))})
			}
			json.NewEncoder(w).Encode(map[string]any{"sha": "abc123", "siblings": sibs})
		case strings.HasPrefix(r.URL.Path, "/org/model/resolve/abc123/"):
			name := strings.TrimPrefix(r.URL.Path, "/org/model/resolve/abc123/")
			c, ok := files[name]
			if !ok {
				http.NotFound(w, r)
				return
			}
			if rg := r.Header.Get("Range"); strings.HasPrefix(rg, "bytes=") {
				from, _ := strconv.Atoi(strings.TrimSuffix(strings.TrimPrefix(rg, "bytes="), "-"))
				if rangeCount != nil {
					atomic.AddInt32(rangeCount, 1)
				}
				w.WriteHeader(206)
				w.Write([]byte(c[from:]))
				return
			}
			w.Write([]byte(c))
		default:
			http.NotFound(w, r)
		}
	}))
}

func withFakeHF(t *testing.T, srv *httptest.Server) (cleanup func(), dir string) {
	old := hfBase
	hfBase = srv.URL
	dir = t.TempDir()
	// point the node config at a temp dir
	cfgDir := t.TempDir()
	_ = cfgDir
	return func() { hfBase = old; srv.Close() }, dir
}

func newTestVLLM(dir string) *vllmRuntime { return newVLLMRuntime() }

func pullTo(t *testing.T, dir, name string, rt *vllmRuntime) error {
	t.Helper()
	// loadNodeConfig reads /etc/wxch/node.json; in tests it falls back to defaults, so override via env-free helper
	cfg := loadNodeConfig()
	cfg.VLLMDir = dir
	return rt.pullWith(context.Background(), cfg, name, func(string, int64, int64) {})
}

func TestPullDownloadsAllFilesAndSkipsDocs(t *testing.T) {
	files := map[string]string{"config.json": `{"a":1}`, "model.safetensors": strings.Repeat("x", 5000), "README.md": "docs", ".gitattributes": "x", "sub/tokenizer.json": "{}"}
	srv := fakeHF(t, files, nil)
	cleanup, dir := withFakeHF(t, srv)
	defer cleanup()
	rt := newTestVLLM(dir)
	if err := pullTo(t, dir, "org/model", rt); err != nil {
		t.Fatal(err)
	}
	for _, f := range []string{"config.json", "model.safetensors", "sub/tokenizer.json"} {
		if _, err := os.Stat(filepath.Join(dir, "org/model", f)); err != nil {
			t.Fatalf("missing %s: %v", f, err)
		}
	}
	for _, f := range []string{"README.md", ".gitattributes"} {
		if _, err := os.Stat(filepath.Join(dir, "org/model", f)); err == nil {
			t.Fatalf("%s should have been skipped", f)
		}
	}
	if b, _ := os.ReadFile(filepath.Join(dir, "org/model", "model.safetensors")); len(b) != 5000 {
		t.Fatalf("size %d", len(b))
	}
}

func TestPullResumesAPartialFile(t *testing.T) {
	big := strings.Repeat("y", 8000)
	files := map[string]string{"config.json": "{}", "model.safetensors": big}
	var ranges int32
	srv := fakeHF(t, files, &ranges)
	cleanup, dir := withFakeHF(t, srv)
	defer cleanup()
	root := filepath.Join(dir, "org/model")
	os.MkdirAll(root, 0755)
	os.WriteFile(filepath.Join(root, "model.safetensors.part"), []byte(big[:3000]), 0644) // an interrupted download
	rt := newTestVLLM(dir)
	if err := pullTo(t, dir, "org/model", rt); err != nil {
		t.Fatal(err)
	}
	if atomic.LoadInt32(&ranges) == 0 {
		t.Fatal("expected a Range request to resume")
	}
	if b, _ := os.ReadFile(filepath.Join(root, "model.safetensors")); string(b) != big {
		t.Fatal("resumed file is corrupt")
	}
}

func TestPullSkipsFilesThatAreAlreadyComplete(t *testing.T) {
	files := map[string]string{"config.json": "{}", "model.safetensors": "abcdef"}
	var ranges int32
	srv := fakeHF(t, files, &ranges)
	cleanup, dir := withFakeHF(t, srv)
	defer cleanup()
	rt := newTestVLLM(dir)
	if err := pullTo(t, dir, "org/model", rt); err != nil {
		t.Fatal(err)
	}
	// second run must not re-download (we count requests via the file content staying identical and no .part left)
	if err := pullTo(t, dir, "org/model", rt); err != nil {
		t.Fatal(err)
	}
	if hasPartial(filepath.Join(dir, "org/model")) {
		t.Fatal("left a .part file behind")
	}
}

func TestPullReplacesOversizedPartialWithoutRange(t *testing.T) {
	files := map[string]string{"config.json": "{}", "model.safetensors": "valid weights"}
	var ranges int32
	srv := fakeHF(t, files, &ranges)
	cleanup, dir := withFakeHF(t, srv)
	defer cleanup()
	root := filepath.Join(dir, "org/model")
	if err := os.MkdirAll(root, 0755); err != nil {
		t.Fatal(err)
	}
	part := filepath.Join(root, "model.safetensors.part")
	if err := os.WriteFile(part, []byte(strings.Repeat("bad", 100)), 0644); err != nil {
		t.Fatal(err)
	}
	if err := pullTo(t, dir, "org/model", newTestVLLM(dir)); err != nil {
		t.Fatal(err)
	}
	if atomic.LoadInt32(&ranges) != 0 {
		t.Fatal("oversized partial must be replaced with a full request")
	}
	got, err := os.ReadFile(filepath.Join(root, "model.safetensors"))
	if err != nil || string(got) != files["model.safetensors"] {
		t.Fatalf("replacement: %q, %v", got, err)
	}
	if _, err := os.Stat(part); !os.IsNotExist(err) {
		t.Fatalf("partial remains: %v", err)
	}
}

func TestPullRejectsConcurrentSameModelAndAllowsRetry(t *testing.T) {
	srv := fakeHF(t, map[string]string{"config.json": "{}", "model.safetensors": "weights"}, nil)
	cleanup, dir := withFakeHF(t, srv)
	defer cleanup()
	rt := newTestVLLM(dir)
	cfg := loadNodeConfig()
	cfg.VLLMDir = dir
	started, release := make(chan struct{}), make(chan struct{})
	result := make(chan error, 1)
	go func() {
		first := true
		result <- rt.pullWith(context.Background(), cfg, "org/model", func(string, int64, int64) {
			if first {
				first = false
				close(started)
				<-release
			}
		})
	}()
	<-started
	err := pullTo(t, dir, "org/model", rt)
	close(release)
	firstErr := <-result
	if err == nil || !strings.Contains(err.Error(), "already downloading") {
		t.Fatalf("duplicate pull: %v", err)
	}
	if firstErr != nil {
		t.Fatal(firstErr)
	}
	if err := pullTo(t, dir, "org/model", rt); err != nil {
		t.Fatalf("retry after completion: %v", err)
	}
}

func TestPullRejectsBadNamesAndMissingRepos(t *testing.T) {
	srv := fakeHF(t, map[string]string{"config.json": "{}"}, nil)
	cleanup, dir := withFakeHF(t, srv)
	defer cleanup()
	rt := newTestVLLM(dir)
	if err := pullTo(t, dir, "../../etc/passwd", rt); err == nil {
		t.Fatal("path traversal name must be rejected")
	}
	if err := pullTo(t, dir, "no/such", rt); err == nil {
		t.Fatal("missing repo must be an error")
	}
	if err := pullTo(t, dir, "justaname", rt); err == nil {
		t.Fatal("name without org must be rejected")
	}
}

func TestPartialModelIsNotListedAsInstalled(t *testing.T) {
	dir := t.TempDir()
	root := filepath.Join(dir, "org/model")
	os.MkdirAll(root, 0755)
	os.WriteFile(filepath.Join(root, "config.json"), []byte("{}"), 0644)
	os.WriteFile(filepath.Join(root, "model.safetensors.part"), []byte("zz"), 0644)
	if !hasPartial(root) {
		t.Fatal("hasPartial should see the .part file")
	}
}

func TestNamingHelpers(t *testing.T) {
	if got := containerOf("nvidia/Qwen3.8-27B-NVFP4"); got != "wxch-vllm-nvidia--qwen3-8-27b-nvfp4" {
		t.Fatal(got)
	}
	if !hfRepoRe.MatchString("nvidia/Qwen3.8-27B-NVFP4") || hfRepoRe.MatchString("a/b/c") || hfRepoRe.MatchString("x") || hfRepoRe.MatchString("a/;rm -rf") {
		t.Fatal("hfRepoRe")
	}
}

func TestPlannedMemory(t *testing.T) {
	cfg := nodeConfig{VLLMKVGB: 8}
	b, u := plannedMemory(20<<30, 128<<30, cfg)
	if b <= 20<<30 || u <= 0 || u > 0.85 {
		t.Fatalf("b=%d u=%f", b, u)
	}
	cfg.VLLMGPUUtil = 0.4
	if _, u := plannedMemory(20<<30, 128<<30, cfg); u != 0.4 {
		t.Fatalf("explicit utilisation must win: %f", u)
	}
	if _, u := plannedMemory(200<<30, 128<<30, nodeConfig{VLLMKVGB: 8}); u != 0.85 {
		t.Fatalf("must be capped: %f", u)
	}
}

// ---- the control layer ----

type fakeCtl struct {
	started  []string
	startErr error
	running  bool
	logs     string
	ports    map[string]int
	srv      map[string]*httptest.Server
	healthy  bool
}

func (f *fakeCtl) Ready(nodeConfig) error { return nil }
func (f *fakeCtl) Containers() ([]vllmContainer, error) {
	var cs []vllmContainer
	for m, p := range f.ports {
		cs = append(cs, vllmContainer{Name: containerOf(m), Model: m, Port: p, State: "running"})
	}
	return cs, nil
}
func (f *fakeCtl) Start(cfg nodeConfig, model string, port int, mem int64, util float64) error {
	f.started = append(f.started, model)
	if f.startErr != nil {
		return f.startErr
	}
	if f.ports == nil {
		f.ports = map[string]int{}
	}
	f.ports[model] = port
	if f.healthy {
		l, err := net.Listen("tcp", "127.0.0.1:"+strconv.Itoa(port))
		if err != nil {
			return err
		}
		srv := &httptest.Server{Listener: l, Config: &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(200) })}}
		srv.Start()
		if f.srv == nil {
			f.srv = map[string]*httptest.Server{}
		}
		f.srv[model] = srv
	}
	return nil
}
func (f *fakeCtl) Running(string) bool     { return f.running }
func (f *fakeCtl) Logs(string, int) string { return f.logs }
func (f *fakeCtl) Stop(model string) error {
	if s := f.srv[model]; s != nil {
		s.Close()
	}
	delete(f.ports, model)
	return nil
}

func withFakeVLLM(t *testing.T, ctl *fakeCtl, models map[string]int64) (*vllmRuntime, func()) {
	dir := t.TempDir()
	for name := range models {
		d := filepath.Join(dir, name)
		os.MkdirAll(d, 0755)
		os.WriteFile(filepath.Join(d, "config.json"), []byte("{}"), 0644)
		os.WriteFile(filepath.Join(d, "model.safetensors"), []byte(strings.Repeat("w", int(models[name]))), 0644)
	}
	cfg := loadNodeConfig()
	cfg.VLLMDir = dir
	cfg.VLLMBasePort = 18100 + int(time.Now().UnixNano()%500)
	cfg.VLLMMaxLen = 4096
	cfgOverride = &cfg
	controlOverride = ctl
	return newVLLMRuntime(), func() {
		cfgOverride, controlOverride = nil, nil
		for _, s := range ctl.srv {
			s.Close()
		}
	}
}

func TestEnsureLoadedStartsTheContainerAndWaitsForHealth(t *testing.T) {
	ctl := &fakeCtl{healthy: true, running: true}
	rt, done := withFakeVLLM(t, ctl, map[string]int64{"nvidia/Qwen3.8-27B-NVFP4": 2000})
	defer done()
	if err := rt.EnsureLoaded(context.Background(), "nvidia/Qwen3.8-27B-NVFP4"); err != nil {
		t.Fatal(err)
	}
	if len(ctl.started) != 1 {
		t.Fatalf("started %v", ctl.started)
	}
	base, err := rt.BaseURL("nvidia/Qwen3.8-27B-NVFP4")
	if err != nil || !strings.HasPrefix(base, "http://127.0.0.1:") {
		t.Fatalf("%q %v", base, err)
	}
	// a second call finds it healthy and does not start another container
	if err := rt.EnsureLoaded(context.Background(), "nvidia/Qwen3.8-27B-NVFP4"); err != nil || len(ctl.started) != 1 {
		t.Fatalf("should reuse the running container: %v %v", err, ctl.started)
	}
}

func TestEnsureLoadedReportsStartFailures(t *testing.T) {
	ctl := &fakeCtl{startErr: fmt.Errorf("docker: no space left")}
	rt, done := withFakeVLLM(t, ctl, map[string]int64{"org/m": 100})
	defer done()
	err := rt.EnsureLoaded(context.Background(), "org/m")
	if err == nil || !strings.Contains(err.Error(), "no space left") {
		t.Fatalf("%v", err)
	}
}

func TestEnsureLoadedShowsTheLogsWhenTheContainerDies(t *testing.T) {
	ctl := &fakeCtl{running: false, logs: "CUDA error: out of memory"}
	rt, done := withFakeVLLM(t, ctl, map[string]int64{"org/m": 100})
	defer done()
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()
	err := rt.EnsureLoaded(ctx, "org/m")
	if err == nil || !strings.Contains(err.Error(), "out of memory") {
		t.Fatalf("the container's last log lines must reach the caller: %v", err)
	}
}

func TestEnsureLoadedRefusesModelsThatAreNotDownloaded(t *testing.T) {
	ctl := &fakeCtl{}
	rt, done := withFakeVLLM(t, ctl, nil)
	defer done()
	if err := rt.EnsureLoaded(context.Background(), "org/missing"); err == nil || len(ctl.started) != 0 {
		t.Fatalf("%v %v", err, ctl.started)
	}
}

// ---- the runner client talks JSON lines over a unix socket ----

func TestRunnerClientRequestsAndResponses(t *testing.T) {
	dir := t.TempDir()
	sock := filepath.Join(dir, "r.sock")
	l, err := net.Listen("unix", sock)
	if err != nil {
		t.Skip("unix sockets unavailable: " + err.Error())
	}
	defer l.Close()
	var got runnerRequest
	go func() {
		for {
			c, err := l.Accept()
			if err != nil {
				return
			}
			line, _ := bufio.NewReader(c).ReadBytes('\n')
			var req runnerRequest
			json.Unmarshal(line, &req)
			got = req
			switch req.Op {
			case "ps":
				c.Write([]byte(`{"ok":true,"data":[{"name":"wxch-vllm-a--b","model":"a/b","port":8100,"mem":123,"state":"running"}]}` + "\n"))
			case "start":
				c.Write([]byte(`{"ok":false,"err":"port 8100 is already used"}` + "\n"))
			default:
				c.Write([]byte(`{"ok":true,"data":true}` + "\n"))
			}
			c.Close()
		}
	}()
	rc := runnerControl{sock: sock}
	cs, err := rc.Containers()
	if err != nil || len(cs) != 1 || cs[0].Model != "a/b" || cs[0].Port != 8100 || cs[0].State != "running" {
		t.Fatalf("%v %v", cs, err)
	}
	cfg := nodeConfig{VLLMMaxLen: 8192, VLLMExtraArgs: []string{"--kv-cache-dtype", "fp8"}}
	err = rc.Start(cfg, "a/b", 8100, 1<<30, 0.3)
	if err == nil || !strings.Contains(err.Error(), "already used") {
		t.Fatalf("runner errors must surface: %v", err)
	}
	if got.Op != "start" || got.Model != "a/b" || got.MaxLen != 8192 || got.Util != 0.3 || len(got.Args) != 2 {
		t.Fatalf("request: %+v", got)
	}
	if !rc.Running("a/b") {
		t.Fatal("running should be true")
	}
	if _, err := (runnerControl{sock: filepath.Join(dir, "nope.sock")}).Containers(); err == nil || !strings.Contains(err.Error(), "runner") {
		t.Fatalf("a missing socket must give a helpful error: %v", err)
	}
}