package proxy_test

// Test scaffolding shared by proxy_test.go and recorder_test.go: the fake health table, the fake
// llama-server upstream, and the rig that builds a whole crossbar over real HTTP.

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"git.wntrmute.dev/kyle/crossbar/internal/config"
	"git.wntrmute.dev/kyle/crossbar/internal/health"
	"git.wntrmute.dev/kyle/crossbar/internal/lease"
	"git.wntrmute.dev/kyle/crossbar/internal/limiter"
	"git.wntrmute.dev/kyle/crossbar/internal/proxy"
	"git.wntrmute.dev/kyle/crossbar/internal/store"
)

// fakeHealth is a hand-set health table that also records MarkDown calls. It lived in the v0
// proxy_test.go; the v1 given test replaces that file, so recorder_test.go (which still exercises
// the nil-lease path through proxy.New) needs it here.
type fakeHealth struct {
	mu     sync.Mutex
	st     map[string]health.Status
	marked []string
}

func (f *fakeHealth) Get(name string) (health.Status, bool) {
	f.mu.Lock()
	defer f.mu.Unlock()
	s, ok := f.st[name]
	return s, ok
}

func (f *fakeHealth) MarkDown(name, reason string) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.marked = append(f.marked, name)
	s := f.st[name]
	s.Healthy = false
	s.LastErr = reason
	f.st[name] = s
}

func (f *fakeHealth) markedHosts() []string {
	f.mu.Lock()
	defer f.mu.Unlock()
	return append([]string{}, f.marked...)
}

// upstream is a llama-server stand-in: streams N chunks with a delay, reports usage/timings in
// the final chunk, counts requests, and can be slowed down or killed.
type upstream struct {
	name  string
	srv   *httptest.Server
	hits  atomic.Int32
	delay time.Duration
	mu    sync.Mutex
	last  recorded
}

type recorded struct{ method, path, host, xff, body string }

func newUpstream(t *testing.T, name string) *upstream {
	u := &upstream{name: name}
	mux := http.NewServeMux()
	mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) { fmt.Fprint(w, `{"status":"ok"}`) })
	mux.HandleFunc("/v1/models", func(w http.ResponseWriter, r *http.Request) {
		u.mu.Lock()
		u.last = recorded{r.Method, r.URL.RequestURI(), r.Host, r.Header.Get("X-Forwarded-For"), ""}
		u.mu.Unlock()
		fmt.Fprint(w, `{"object":"list","data":[{"id":"shared"},{"id":"`+name+`-only"}]}`)
	})
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		u.hits.Add(1)
		b, _ := io.ReadAll(r.Body)
		u.mu.Lock()
		u.last = recorded{r.Method, r.URL.RequestURI(), r.Host, r.Header.Get("X-Forwarded-For"), string(b)}
		u.mu.Unlock()
		var req struct {
			Stream bool `json:"stream"`
		}
		_ = json.Unmarshal(b, &req)
		w.Header().Set("X-Upstream", name)
		time.Sleep(u.delay)
		if !req.Stream {
			w.Header().Set("Content-Type", "application/json")
			fmt.Fprintf(w, `{"choices":[{"message":{"role":"assistant","content":"hi from %s"}}],"usage":{"prompt_tokens":100,"completion_tokens":10,"total_tokens":110},"timings":{"prompt_n":100,"cache_n":90,"predicted_n":10,"predicted_ms":50.0}}`, name)
			return
		}
		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(200)
		fl := w.(http.Flusher)
		for i := 0; i < 3; i++ {
			fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"content\":\"%s %d \"}}]}\n\n", name, i)
			fl.Flush()
			time.Sleep(10 * time.Millisecond)
		}
		fmt.Fprint(w, `data: {"choices":[],"usage":{"prompt_tokens":200,"completion_tokens":20,"total_tokens":220},"timings":{"prompt_n":200,"cache_n":150,"predicted_n":20,"predicted_ms":80.0}}`+"\n\n")
		fl.Flush()
		fmt.Fprint(w, "data: [DONE]\n\n")
	})
	u.srv = httptest.NewServer(mux)
	t.Cleanup(u.srv.Close)
	return u
}

func (u *upstream) lastReq() recorded { u.mu.Lock(); defer u.mu.Unlock(); return u.last }

// rig is one crossbar: config, real health table (polled once), real lease table over a real
// SQLite store, real limiter, the proxy handler served by httptest.
type rig struct {
	t      *testing.T
	cfg    *config.Config
	health *health.Table
	store  *store.Store
	leases *lease.Table
	lim    *limiter.Limiter
	front  *httptest.Server
}

// newRig builds crossbar from a config text where %s placeholders are the upstream base URLs.
func newRig(t *testing.T, cfgText string, ups ...*upstream) *rig {
	urls := make([]any, len(ups))
	for i, u := range ups {
		urls[i] = u.srv.URL
	}
	cfg, err := config.Parse(strings.NewReader(fmt.Sprintf(cfgText, urls...)))
	if err != nil {
		t.Fatal(err)
	}
	bases := map[string]string{}
	for name, h := range cfg.Hosts {
		bases[name] = h.BaseURL
	}
	ht := health.New(bases, time.Hour, nil)
	ht.PollOnce(t.Context())
	st, err := store.Open(filepath.Join(t.TempDir(), "crossbar.db"))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = st.Close() })
	lim := limiter.New()
	for name, h := range cfg.Hosts {
		for model, m := range h.Models {
			lim.Configure(name, model, m.Parallel, cfg.QueueMax)
		}
	}
	lt, err := lease.New(st, proxy.HostView(ht, cfg), proxy.Chooser(cfg, ht, lim), cfg.LeaseIdle.Duration)
	if err != nil {
		t.Fatal(err)
	}
	p := proxy.New(cfg, ht, lt, lim, st, nil)
	front := httptest.NewServer(p)
	t.Cleanup(front.Close)
	return &rig{t: t, cfg: cfg, health: ht, store: st, leases: lt, lim: lim, front: front}
}

const twoHosts = `
listen = "127.0.0.1:1"
queue_max = 1
lease_idle = "30m"
[hosts.alpha]
base_url = %q
weight = 1.0
models = { "shared" = { parallel = 2 }, "alpha-only" = { } }
[hosts.beta]
base_url = %q
weight = 2.0
models = { "shared" = { parallel = 2 }, "beta-only" = { } }
[routes.r]
hosts = ["alpha", "beta"]
default_model = "shared"
[routes.other]
hosts = ["alpha"]
`

func conversation(id, turn int) string {
	msgs := fmt.Sprintf(`{"role":"system","content":"project"},{"role":"user","content":"conversation %d opening"}`, id)
	for i := 1; i < turn; i++ {
		msgs += fmt.Sprintf(`,{"role":"assistant","content":"ok"},{"role":"user","content":"turn %d"}`, i)
	}
	return `{"model":"shared","stream":false,"messages":[` + msgs + `]}`
}

func (r *rig) post(path, body string, hdr ...string) *http.Response {
	req, _ := http.NewRequest(http.MethodPost, r.front.URL+path, strings.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	for i := 0; i+1 < len(hdr); i += 2 {
		req.Header.Set(hdr[i], hdr[i+1])
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		r.t.Fatal(err)
	}
	return resp
}

func drain(resp *http.Response) string {
	b, _ := io.ReadAll(resp.Body)
	resp.Body.Close()
	return string(b)
}