package main

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
)

// fakeLM imitates the parts of LM Studio's REST API that the runtime uses.
type fakeLM struct {
	mu       sync.Mutex
	loaded   map[string]string // key -> instance id
	polls    int
	loadBody map[string]any
}

func (f *fakeLM) handler() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		f.mu.Lock()
		defer f.mu.Unlock()
		switch {
		case r.URL.Path == "/api/v1/models":
			type inst struct {
				ID string `json:"id"`
			}
			mk := func(key string, size int64) map[string]any {
				var li []inst
				if id, ok := f.loaded[key]; ok {
					li = append(li, inst{id})
				}
				return map[string]any{"type": "llm", "key": key, "size_bytes": size, "loaded_instances": li}
			}
			json.NewEncoder(w).Encode(map[string]any{"models": []any{mk("nvidia/nemotron-omni", 20e9), mk("org/small", 2e9), map[string]any{"type": "embedding", "key": "emb/x", "size_bytes": 1e8}}})
		case r.URL.Path == "/api/v1/models/load" && r.Method == "POST":
			json.NewDecoder(r.Body).Decode(&f.loadBody)
			key, _ := f.loadBody["model"].(string)
			f.loaded[key] = key + "-inst"
			json.NewEncoder(w).Encode(map[string]any{"type": "llm", "instance_id": key + "-inst", "status": "loaded"})
		case r.URL.Path == "/api/v1/models/unload" && r.Method == "POST":
			var b map[string]string
			json.NewDecoder(r.Body).Decode(&b)
			for k, id := range f.loaded {
				if id == b["instance_id"] {
					delete(f.loaded, k)
				}
			}
			json.NewEncoder(w).Encode(map[string]string{"instance_id": b["instance_id"]})
		case r.URL.Path == "/api/v1/models/download" && r.Method == "POST":
			var b map[string]any
			json.NewDecoder(r.Body).Decode(&b)
			if m, _ := b["model"].(string); strings.Contains(m, "already") {
				json.NewEncoder(w).Encode(map[string]any{"status": "already_downloaded"})
				return
			}
			f.loadBody = b
			json.NewEncoder(w).Encode(map[string]any{"job_id": "job1", "status": "downloading", "total_size_bytes": 1000})
		case strings.HasPrefix(r.URL.Path, "/api/v1/models/download/status/"):
			f.polls++
			if f.polls < 3 {
				json.NewEncoder(w).Encode(map[string]any{"status": "downloading", "total_size_bytes": 1000, "downloaded_bytes": f.polls * 300})
			} else {
				json.NewEncoder(w).Encode(map[string]any{"status": "completed", "total_size_bytes": 1000, "downloaded_bytes": 1000})
			}
		default:
			http.NotFound(w, r)
		}
	})
}

func useFakeLM(t *testing.T) (*fakeLM, func()) {
	f := &fakeLM{loaded: map[string]string{}}
	srv := httptest.NewServer(f.handler())
	t.Setenv("WXCH_TEST_LMURL", srv.URL)
	lmTestURL = srv.URL
	return f, func() { lmTestURL = ""; srv.Close() }
}

func TestLMStudioListsOnlyLLMs(t *testing.T) {
	_, done := useFakeLM(t)
	defer done()
	rt := newLMStudioRuntime()
	inst, err := rt.Installed()
	if err != nil || len(inst) != 2 || inst["nvidia/nemotron-omni"] != 20e9 {
		t.Fatalf("%v %v", inst, err)
	}
	if rt.Ready() != nil {
		t.Fatal("should be ready")
	}
}

func TestLMStudioLoadsWithOurContextAndUnloads(t *testing.T) {
	f, done := useFakeLM(t)
	defer done()
	rt := newLMStudioRuntime()
	if err := rt.EnsureLoaded(context.Background(), "org/small"); err != nil {
		t.Fatal(err)
	}
	if f.loadBody["context_length"] == nil || f.loadBody["model"] != "org/small" {
		t.Fatalf("load request: %v", f.loadBody)
	}
	l, _ := rt.Loaded()
	if _, ok := l["org/small"]; !ok {
		t.Fatalf("should be loaded: %v", l)
	}
	if err := rt.EnsureLoaded(context.Background(), "org/small"); err != nil { // already loaded: no second load
		t.Fatal(err)
	}
	if err := rt.Unload("org/small"); err != nil {
		t.Fatal(err)
	}
	if l, _ := rt.Loaded(); len(l) != 0 {
		t.Fatalf("still loaded: %v", l)
	}
	if err := rt.Unload("nope/none"); err == nil {
		t.Fatal("unknown model must be an error")
	}
}

func TestLMStudioPullPollsUntilComplete(t *testing.T) {
	f, done := useFakeLM(t)
	defer done()
	rt := newLMStudioRuntime()
	var seen []int64
	err := rt.Pull(context.Background(), "ibm/granite-4-micro", func(st string, total, d int64) { seen = append(seen, d) })
	if err != nil || len(seen) < 3 || seen[len(seen)-1] != 1000 {
		t.Fatalf("progress %v err %v", seen, err)
	}
	_ = f
}

func TestLMStudioPullHuggingFaceWithQuantization(t *testing.T) {
	f, done := useFakeLM(t)
	defer done()
	rt := newLMStudioRuntime()
	if err := rt.Pull(context.Background(), "lmstudio-community/Foo-GGUF@Q4_K_M", func(string, int64, int64) {}); err != nil {
		t.Fatal(err)
	}
	if f.loadBody["quantization"] != "Q4_K_M" || f.loadBody["model"] != "https://huggingface.co/lmstudio-community/Foo-GGUF" {
		t.Fatalf("download request: %v", f.loadBody)
	}
	if err := rt.Pull(context.Background(), "already-there", func(string, int64, int64) {}); err != nil {
		t.Fatalf("already_downloaded must count as success: %v", err)
	}
}

func TestLMStudioNotRunningIsReportedClearly(t *testing.T) {
	lmTestURL = "http://127.0.0.1:1" // nothing listens here
	defer func() { lmTestURL = "" }()
	rt := newLMStudioRuntime()
	err := rt.Ready()
	if err == nil || !strings.Contains(err.Error(), "LM Studio") {
		t.Fatalf("%v", err)
	}
}

func TestLMStudioDeleteRemovesTheModelFolderAndStaysInsideTheRoot(t *testing.T) {
	_, done := useFakeLM(t)
	defer done()
	root := t.TempDir()
	dir := filepath.Join(root, "org", "Small-GGUF")
	os.MkdirAll(dir, 0755)
	os.WriteFile(filepath.Join(dir, "Small-Q4.gguf"), []byte("w"), 0644)
	os.WriteFile(filepath.Join(dir, "mmproj-Small.gguf"), []byte("v"), 0644)
	other := filepath.Join(root, "org", "Keep-GGUF")
	os.MkdirAll(other, 0755)
	os.WriteFile(filepath.Join(other, "Keep.gguf"), []byte("k"), 0644)

	t.Setenv("WXCH_TEST_LMROOT", root)
	oldList := lmListFiles
	lmListFiles = func(nodeConfig) ([]lmFile, error) {
		return []lmFile{{"org/small", "org/Small-GGUF/Small-Q4.gguf"}, {"org/keep", "org/Keep-GGUF/Keep.gguf"}, {"evil", "../../etc/passwd"}}, nil
	}
	oldRoot, oldWait := lmRootOverride, lmDeleteWait
	lmRootOverride, lmDeleteWait = root, 0
	defer func() { lmListFiles, lmRootOverride, lmDeleteWait = oldList, oldRoot, oldWait }()

	rt := newLMStudioRuntime()
	if err := rt.Delete("org/small"); err != nil {
		t.Fatal(err)
	}
	if _, err := os.Stat(dir); err == nil {
		t.Fatal("the model folder (including the leftover mmproj) should be gone")
	}
	if _, err := os.Stat(filepath.Join(other, "Keep.gguf")); err != nil {
		t.Fatal("another model was touched")
	}
	if err := rt.Delete("evil"); err == nil {
		t.Fatal("a path outside the models folder must be refused")
	}
	if err := rt.Delete("nope"); err == nil {
		t.Fatal("unknown model must be an error")
	}
}

func TestLMStudioDeleteRefusesALoadedModel(t *testing.T) {
	f, done := useFakeLM(t)
	defer done()
	f.loaded["org/small"] = "org/small-inst"
	rt := newLMStudioRuntime()
	if err := rt.Delete("org/small"); err == nil || !strings.Contains(err.Error(), "unload") {
		t.Fatalf("%v", err)
	}
}