package cli

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"slices"
	"sort"
	"strings"
	"sync"
	"syscall"
	"text/tabwriter"
	"time"

	"github.com/spf13/cobra"

	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/guard"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router"
	storepkg "gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// What llmgw adds to a request, measured: the same load straight at an
// engine that answers at once, and through the gateway in front of it —
// bare, with authentication and usage accounting, and with the guardrails'
// pattern checks besides. The engine's own time is the same on both paths,
// so what differs is the gateway's.
//
// Everything runs in this process, on loopback: no cluster, no GPU, and the
// numbers are this machine's. On a machine shared with other work (a VM on
// an oversubscribed host) they move with it; compare modes run together,
// not figures from different runs.

// overheadMode is one way a request reaches the engine.
type overheadMode struct {
	name        string
	gateway     bool
	auth, guard bool
}

var overheadModes = []overheadMode{
	{name: "direct"},
	{name: "gateway", gateway: true},
	{name: "gateway+auth", gateway: true, auth: true},
	{name: "gateway+auth+guard", gateway: true, auth: true, guard: true},
}

// overheadStep is one mode at one request rate.
type overheadStep struct {
	Mode     string        `json:"mode"`
	Rate     int           `json:"rate"`
	Sent     int           `json:"sent"`
	Failed   int           `json:"failed"`
	P50      time.Duration `json:"p50"`
	P99      time.Duration `json:"p99"`
	AddedP50 time.Duration `json:"added_p50"`
	AddedP99 time.Duration `json:"added_p99"`
	// CPUPer1k is the process's CPU time per thousand requests beyond what
	// the direct mode spent at the same rate: the gateway's.
	CPUPer1k time.Duration `json:"cpu_per_1k"`
	cpu      time.Duration
}

func newBenchOverheadCmd(asJSON *bool) *cobra.Command {
	var (
		rates []int
		each  time.Duration
	)
	cmd := &cobra.Command{
		Use:   "overhead",
		Short: "Measure what the gateway adds to a request, against an engine that answers at once",
		Long: "Runs an engine that answers at once and a gateway in front of it, in this\n" +
			"process on loopback, and sends the same load at each rate straight to the\n" +
			"engine and through the gateway: bare, with authentication and usage\n" +
			"accounting, and with the guardrails' pattern checks too. Reports each\n" +
			"request's latency at p50 and p99, what the gateway added to it, and the\n" +
			"gateway's CPU time per thousand requests.\n\n" +
			"Needs no cluster. The figures are this machine's: on one shared with other\n" +
			"work they move with it, so compare the modes of one run, and note the\n" +
			"machine's load beside any figure you keep.",
		Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			steps, err := runOverhead(cmd.Context(), rates, each)
			if err != nil {
				return err
			}
			if *asJSON {
				return json.NewEncoder(cmd.OutOrStdout()).Encode(steps)
			}
			printOverhead(cmd.OutOrStdout(), steps)
			return nil
		},
	}
	f := cmd.Flags()
	f.IntSliceVar(&rates, "rates", []int{100, 500, 1000, 2000}, "requests a second to send, one step each")
	f.DurationVar(&each, "each", 3*time.Second, "how long each step sends for")
	return cmd
}

func printOverhead(w io.Writer, steps []overheadStep) {
	tw := tabwriter.NewWriter(w, 0, 0, 2, ' ', 0)
	fmt.Fprintln(tw, "MODE\tRATE/S\tSENT\tFAILED\tP50\tP99\t+P50\t+P99\tCPU/1K REQ")
	for _, s := range steps {
		added := func(d time.Duration) string {
			if s.Mode == "direct" {
				return "-"
			}
			return msOf(d)
		}
		cpu := "-"
		if s.Mode != "direct" {
			cpu = msOf(s.CPUPer1k)
		}
		fmt.Fprintf(tw, "%s\t%d\t%d\t%d\t%s\t%s\t%s\t%s\t%s\n", s.Mode, s.Rate, s.Sent, s.Failed,
			msOf(s.P50), msOf(s.P99), added(s.AddedP50), added(s.AddedP99), cpu)
	}
	_ = tw.Flush()
	fmt.Fprintf(w, "\n%d CPUs, %s. Loopback, in one process: the figures are this machine's.\n",
		runtime.NumCPU(), loadAverage())
}

func msOf(d time.Duration) string { return fmt.Sprintf("%.2fms", float64(d)/float64(time.Millisecond)) }

// loadAverage is the machine's, where it can be read, to keep beside the
// figures.
func loadAverage() string {
	if b, err := os.ReadFile("/proc/loadavg"); err == nil {
		if f := strings.Fields(string(b)); len(f) >= 3 {
			return "load " + strings.Join(f[:3], " ")
		}
	}
	// macOS: "{ 1.23 1.10 0.98 }".
	if b, err := exec.Command("sysctl", "-n", "vm.loadavg").Output(); err == nil {
		if f := strings.Fields(strings.Trim(strings.TrimSpace(string(b)), "{}")); len(f) >= 3 {
			return "load " + strings.Join(f[:3], " ")
		}
	}
	return "load not read"
}

// runOverhead measures every mode at every rate.
func runOverhead(ctx context.Context, rates []int, each time.Duration) ([]overheadStep, error) {
	engine, err := instantEngine()
	if err != nil {
		return nil, err
	}
	defer engine.Close()
	dir, err := os.MkdirTemp("", "llmgw-overhead-")
	if err != nil {
		return nil, err
	}
	defer os.RemoveAll(dir)

	type target struct {
		url, key string
		stop     func()
	}
	targets := map[string]target{}
	defer func() {
		for _, t := range targets {
			if t.stop != nil {
				t.stop()
			}
		}
	}()
	for _, m := range overheadModes {
		if !m.gateway {
			targets[m.name] = target{url: engine.URL}
			continue
		}
		url, key, stop, err := overheadGateway(ctx, engine.URL, filepath.Join(dir, m.name+".db"), m)
		if err != nil {
			return nil, fmt.Errorf("%s: %w", m.name, err)
		}
		targets[m.name] = target{url: url, key: key, stop: stop}
	}

	client := &http.Client{Transport: &http.Transport{MaxIdleConns: 1024, MaxIdleConnsPerHost: 1024}, Timeout: 30 * time.Second}
	var out []overheadStep
	for _, rate := range rates {
		var direct overheadStep
		for _, m := range overheadModes {
			t := targets[m.name]
			s := loadStep(ctx, client, t.url, t.key, rate, each)
			s.Mode = m.name
			if m.name == "direct" {
				direct = s
			} else {
				s.AddedP50, s.AddedP99 = max(0, s.P50-direct.P50), max(0, s.P99-direct.P99)
				if s.Sent > 0 {
					s.CPUPer1k = max(0, s.cpu-direct.cpu) * 1000 / time.Duration(s.Sent)
				}
			}
			out = append(out, s)
		}
	}
	return out, nil
}

// loadStep sends rate requests a second for d, each in its own goroutine,
// and measures each one's latency and the process's CPU meanwhile.
func loadStep(ctx context.Context, client *http.Client, url, key string, rate int, d time.Duration) overheadStep {
	s := overheadStep{Rate: rate}
	body := []byte(`{"model":"bench","max_tokens":8,"messages":[{"role":"user","content":"What is the capital of France?"}]}`)
	var (
		mu    sync.Mutex
		lat   []time.Duration
		wg    sync.WaitGroup
		total = int(float64(rate) * d.Seconds())
	)
	interval := time.Second / time.Duration(max(rate, 1))
	cpu0 := processCPU()
	start := time.Now()
	for i := range total {
		if ctx.Err() != nil {
			break
		}
		if wait := time.Until(start.Add(time.Duration(i) * interval)); wait > 0 {
			time.Sleep(wait)
		}
		wg.Add(1)
		go func() {
			defer wg.Done()
			req, _ := http.NewRequestWithContext(ctx, http.MethodPost, url+"/v1/chat/completions", bytes.NewReader(body))
			req.Header.Set("Content-Type", "application/json")
			if key != "" {
				req.Header.Set("Authorization", "Bearer "+key)
			}
			t := time.Now()
			resp, err := client.Do(req)
			ok := err == nil && resp.StatusCode == http.StatusOK
			if resp != nil {
				_, _ = io.Copy(io.Discard, resp.Body)
				resp.Body.Close()
			}
			el := time.Since(t)
			mu.Lock()
			s.Sent++
			if ok {
				lat = append(lat, el)
			} else {
				s.Failed++
			}
			mu.Unlock()
		}()
	}
	wg.Wait()
	s.cpu = processCPU() - cpu0
	s.P50, s.P99 = percentile(lat, 0.50), percentile(lat, 0.99)
	return s
}

func percentile(ds []time.Duration, p float64) time.Duration {
	if len(ds) == 0 {
		return 0
	}
	ds = slices.Clone(ds)
	sort.Slice(ds, func(i, j int) bool { return ds[i] < ds[j] })
	i := int(p * float64(len(ds)-1))
	return ds[i]
}

// processCPU is the CPU time this process has used, user and system.
func processCPU() time.Duration {
	var ru syscall.Rusage
	if err := syscall.Getrusage(syscall.RUSAGE_SELF, &ru); err != nil {
		return 0
	}
	return time.Duration(ru.Utime.Nano() + ru.Stime.Nano())
}

// instantEngine is an OpenAI-compatible engine that answers at once, with
// the metrics the router scrapes.
func instantEngine() (*httpServer, error) {
	mux := http.NewServeMux()
	mux.HandleFunc("GET /metrics", func(w http.ResponseWriter, _ *http.Request) {
		fmt.Fprint(w, "# TYPE vllm:num_requests_running gauge\nvllm:num_requests_running 0\n"+
			"# TYPE vllm:num_requests_waiting gauge\nvllm:num_requests_waiting 0\n"+
			"# TYPE vllm:kv_cache_usage_perc gauge\nvllm:kv_cache_usage_perc 0.1\n")
	})
	mux.HandleFunc("GET /health", func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) })
	mux.HandleFunc("POST /v1/chat/completions", func(w http.ResponseWriter, r *http.Request) {
		_, _ = io.Copy(io.Discard, r.Body)
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprint(w, `{"id":"b","object":"chat.completion","model":"bench","choices":[{"index":0,"message":{"role":"assistant","content":"Paris."},"finish_reason":"stop"}],"usage":{"prompt_tokens":14,"completion_tokens":2,"total_tokens":16}}`)
	})
	return serveLoopback(mux)
}

// httpServer is a handler served on a loopback port.
type httpServer struct {
	URL string
	srv *http.Server
}

func (s *httpServer) Close() { _ = s.srv.Close() }

func serveLoopback(h http.Handler) (*httpServer, error) {
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return nil, err
	}
	srv := &http.Server{Handler: h, ReadHeaderTimeout: 10 * time.Second}
	go func() { _ = srv.Serve(ln) }()
	return &httpServer{URL: "http://" + ln.Addr().String(), srv: srv}, nil
}

// overheadGateway starts a gateway in front of the engine as the mode asks,
// and returns its address and, with auth, a key.
func overheadGateway(ctx context.Context, engineURL, storePath string, m overheadMode) (string, string, func(), error) {
	cfg := config.Default()
	cfg.Listen = "127.0.0.1:0"
	cfg.Router.ScrapeInterval = 200 * time.Millisecond
	cfg.Endpoints = []config.Endpoint{{ID: "bench", URL: engineURL, Models: []string{"bench"}}}
	key := ""
	if m.auth {
		var err error
		if key, err = overheadKey(ctx, storePath); err != nil {
			return "", "", nil, err
		}
		cfg.Auth.Enabled, cfg.Auth.Store = true, storePath
		// Limits would measure themselves refusing: off.
		cfg.Auth.DefaultRPM, cfg.Auth.DefaultTPM = 0, 0
	}
	rt, err := router.New(cfg, slog.New(slog.DiscardHandler))
	if err != nil {
		return "", "", nil, err
	}
	if m.guard {
		g, err := baselineGuard()
		if err != nil {
			return "", "", nil, err
		}
		rt.SetGuard(g)
	}
	rctx, cancel := context.WithCancel(ctx)
	rt.Start(rctx)
	srv, err := serveLoopback(rt.Handler())
	if err != nil {
		cancel()
		return "", "", nil, err
	}
	stop := func() { srv.Close(); cancel(); rt.Stop() }
	// Ready once a request goes through: the endpoint is scraped first.
	deadline := time.Now().Add(10 * time.Second)
	for {
		body := `{"model":"bench","messages":[{"role":"user","content":"hi"}]}`
		req, _ := http.NewRequestWithContext(ctx, http.MethodPost, srv.URL+"/v1/chat/completions", strings.NewReader(body))
		req.Header.Set("Content-Type", "application/json")
		if key != "" {
			req.Header.Set("Authorization", "Bearer "+key)
		}
		resp, err := http.DefaultClient.Do(req)
		if err == nil {
			_, _ = io.Copy(io.Discard, resp.Body)
			resp.Body.Close()
			if resp.StatusCode == http.StatusOK {
				return srv.URL, key, stop, nil
			}
		}
		if time.Now().After(deadline) {
			stop()
			status := "no answer"
			if err == nil {
				status = resp.Status
			}
			return "", "", nil, fmt.Errorf("the gateway did not serve within 10s (%s)", status)
		}
		time.Sleep(100 * time.Millisecond)
	}
}

// overheadKey makes a store with one project and a key in it.
func overheadKey(ctx context.Context, path string) (string, error) {
	st, err := storepkg.Open(ctx, path)
	if err != nil {
		return "", err
	}
	defer st.Close()
	org, err := st.CreateOrg(ctx, "bench")
	if err != nil {
		return "", err
	}
	team, err := st.CreateTeam(ctx, org, "bench")
	if err != nil {
		return "", err
	}
	project, err := st.CreateProject(ctx, team, "bench")
	if err != nil {
		return "", err
	}
	secret, k, err := storepkg.NewKey("key-bench", "bench", project)
	if err != nil {
		return "", err
	}
	k.CreatedAt = time.Now()
	return secret, st.CreateKey(ctx, k)
}

// baselineGuard is the policy init installs, without its judge: the pattern
// checks every request passes.
func baselineGuard() (guard.Guard, error) {
	m, err := config.Parse([]byte(guardSection(false)))
	if err != nil {
		return nil, err
	}
	cfg, err := config.FromMap(m)
	if err != nil {
		return nil, err
	}
	return guard.NewGate(cfg.Guard, nil, slog.New(slog.DiscardHandler))
}