// Run inside the cluster: Pod IPs are not directly reachable from macOS.
package main

import (
	"bytes"
	"context"
	"crypto/subtle"
	"encoding/json"
	"fmt"
	"github.com/redis/go-redis/v9"
	"io"
	dis "k8s.io/api/discovery/v1"
	meta "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
	"log"
	"math"
	"net"
	"net/http"
	"os"
	"os/signal"
	"sort"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"
)

type Tenant struct {
	Token  string
	Models []string
}
type Router struct {
	sync.Mutex
	routes   map[string][]string
	active   map[string]int
	windows  map[string][]time.Time
	external *redis.Client
	penalty  map[string]float64
	users    map[string]Tenant
	http     *http.Client
	last     time.Time
}

func (g *Router) refresh(ctx context.Context, k kubernetes.Interface) error {
	list, e := k.DiscoveryV1().EndpointSlices("modeldock-lab").List(ctx, meta.ListOptions{})
	if e != nil {
		return e
	}
	out := map[string][]string{}
	for _, s := range list.Items {
		model := s.Labels[dis.LabelServiceName]
		for _, p := range s.Ports {
			if p.Port == nil {
				continue
			}
			for _, ep := range s.Endpoints {
				if ep.Conditions.Ready == nil || !*ep.Conditions.Ready || (ep.Conditions.Terminating != nil && *ep.Conditions.Terminating) {
					continue
				}
				for _, ip := range ep.Addresses {
					out[model] = append(out[model], "http://"+net.JoinHostPort(ip, fmt.Sprint(*p.Port)))
				}
			}
		}
	}
	// Probe at most a bounded number per refresh; absent metrics fall back to local inflight.
	penalties := map[string]float64{}
	var wg sync.WaitGroup
	var m sync.Mutex
	budget, cancel := context.WithTimeout(ctx, 2*time.Second)
	defer cancel()
	jobs := make(chan string, 16)
	probe := func(p string) {
		q, _ := http.NewRequestWithContext(budget, "GET", p+"/metrics", nil)
		res, er := g.http.Do(q)
		if er != nil {
			return
		}
		defer res.Body.Close()
		if res.StatusCode != 200 {
			return
		}
		b, er := io.ReadAll(io.LimitReader(res.Body, 1<<20))
		if er == nil {
			m.Lock()
			penalties[p] = score(string(b))
			m.Unlock()
		}
	}
	for range 8 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for p := range jobs {
				if budget.Err() != nil {
					return
				}
				probe(p)
			}
		}()
	}
enqueue:
	for _, ps := range out {
		for _, p := range ps {
			select {
			case jobs <- p:
			case <-budget.Done():
				break enqueue
			}
		}
	}
	close(jobs)
	wg.Wait()
	g.Lock()
	defer g.Unlock()
	g.penalty = penalties
	g.routes = out
	g.last = time.Now()
	return nil
}
func score(metrics string) float64 {
	var pressure float64
	for _, line := range strings.Split(metrics, "\n") {
		f := strings.Fields(line)
		if len(f) < 2 || f[0] == "#" {
			continue
		}
		v, e := strconv.ParseFloat(f[1], 64)
		if e != nil || v < 0 || math.IsNaN(v) || math.IsInf(v, 0) {
			continue
		}
		switch strings.Split(f[0], "{")[0] {
		case "modeldock_waiting_requests", "vllm:num_requests_waiting":
			pressure += 2 * v
		case "modeldock_running_requests", "vllm:num_requests_running":
			pressure += v
		case "vllm:kv_cache_usage_perc", "vllm:gpu_cache_usage_perc":
			pressure += 10 * v
		}
	}
	return pressure
}

var budgetScript = redis.NewScript("local n=redis.call('INCR',KEYS[1]); if n==1 then redis.call('PEXPIRE',KEYS[1],60000) end; return n")

func (g *Router) serve(w http.ResponseWriter, r *http.Request) {
	key, ok := strings.CutPrefix(r.Header.Get("Authorization"), "Bearer ")
	if !ok {
		http.Error(w, "bearer required", 401)
		return
	}
	var tenant string
	for n, u := range g.users {
		if u.Token != "" && subtle.ConstantTimeCompare([]byte(key), []byte(u.Token)) == 1 {
			tenant = n
			break
		}
	}
	if tenant == "" {
		http.Error(w, "unauthorized", 401)
		return
	}
	body, e := io.ReadAll(http.MaxBytesReader(w, r.Body, 65536))
	var b struct{ Model string }
	if e != nil || json.Unmarshal(body, &b) != nil {
		http.Error(w, "invalid request", 400)
		return
	}
	allowed := false
	for _, m := range g.users[tenant].Models {
		if b.Model == m {
			allowed = true
		}
	}
	if !allowed {
		http.Error(w, "model outside tenant scope", 403)
		return
	}
	if g.external != nil {
		n, e := budgetScript.Run(r.Context(), g.external, []string{"modeldock:budget:" + tenant}).Int64()
		if e != nil {
			http.Error(w, "budget unavailable", 503)
			return
		}
		if n > 60 {
			http.Error(w, "shared rate limit", 429)
			return
		}
	}
	g.Lock()
	now := time.Now()
	window := g.windows[tenant]
	for len(window) > 0 && now.Sub(window[0]) >= time.Minute {
		window = window[1:]
	}
	if len(window) >= 60 {
		g.Unlock()
		http.Error(w, "rate limit", 429)
		return
	}
	g.windows[tenant] = append(window, now)
	peers := append([]string{}, g.routes[b.Model]...)
	if now.Sub(g.last) > 15*time.Second {
		peers = nil
	}
	g.Unlock()
	for len(peers) > 0 {
		g.Lock()
		sort.Slice(peers, func(i, j int) bool {
			return float64(g.active[peers[i]])+g.penalty[peers[i]] < float64(g.active[peers[j]])+g.penalty[peers[j]]
		})
		p := peers[0]
		peers = peers[1:]
		g.active[p]++
		g.Unlock()
		req, _ := http.NewRequestWithContext(r.Context(), "POST", p+"/v1/chat/completions", bytes.NewReader(body))
		req.Header.Set("Content-Type", "application/json")
		resp, err := g.http.Do(req)
		if err != nil || resp.StatusCode == 503 {
			if resp != nil {
				resp.Body.Close()
			}
			g.Lock()
			g.active[p]--
			g.Unlock()
			if r.Context().Err() != nil {
				return
			}
			continue
		}
		defer resp.Body.Close()
		defer func() { g.Lock(); g.active[p]--; g.Unlock() }()
		w.Header().Set("Content-Type", resp.Header.Get("Content-Type"))
		w.Header().Set("Cache-Control", "no-store")
		w.WriteHeader(resp.StatusCode)
		rc := http.NewResponseController(w)
		buf := make([]byte, 8192)
		for {
			n, e := resp.Body.Read(buf)
			if n > 0 {
				rc.SetWriteDeadline(time.Now().Add(10 * time.Second))
				if _, er := w.Write(buf[:n]); er != nil {
					return
				}
				if rc.Flush() != nil {
					return
				}
			}
			if e != nil {
				return
			}
		}
	}
	http.Error(w, "no available endpoint", 503)
}
func main() {
	cfg, e := rest.InClusterConfig()
	if e != nil {
		panic(e)
	}
	k, e := kubernetes.NewForConfig(cfg)
	if e != nil {
		panic(e)
	}
	g := Router{active: map[string]int{}, windows: map[string][]time.Time{}, http: &http.Client{Timeout: 90 * time.Second, Transport: &http.Transport{DialContext: (&net.Dialer{Timeout: 2 * time.Second}).DialContext, MaxIdleConnsPerHost: 8, ResponseHeaderTimeout: 15 * time.Second}}}
	if uri := os.Getenv("REDIS_URL"); uri != "" {
		o, err := redis.ParseURL(uri)
		if err != nil {
			panic(err)
		}
		g.external = redis.NewClient(o)
		defer g.external.Close()
	}
	if e = json.Unmarshal([]byte(os.Getenv("TENANT_KEYS_JSON")), &g.users); e != nil {
		panic(e)
	}
	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGTERM, os.Interrupt)
	defer stop()
	go func() {
		for {
			if e := g.refresh(ctx, k); e != nil {
				log.Println(e)
			}
			select {
			case <-ctx.Done():
				return
			case <-time.After(3 * time.Second):
			}
		}
	}()
	mux := http.NewServeMux()
	mux.HandleFunc("POST /v1/chat/completions", g.serve)
	mux.HandleFunc("GET /ready", func(w http.ResponseWriter, r *http.Request) {
		g.Lock()
		defer g.Unlock()
		if time.Since(g.last) > 15*time.Second {
			http.Error(w, "discovery stale", 503)
		}
	})
	srv := http.Server{Addr: ":8080", Handler: mux, ReadHeaderTimeout: 5 * time.Second}
	drained := make(chan struct{})
	go func() {
		defer close(drained)
		<-ctx.Done()
		c, cancel := context.WithTimeout(context.Background(), 12*time.Second)
		defer cancel()
		srv.Shutdown(c)
	}()
	if e = srv.ListenAndServe(); e != http.ErrServerClosed {
		panic(e)
	}
	<-drained
}