package main

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"io"
	"net/http"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"
)

// OpenAI-compatible gateway served on api.wxch.org. A request is authenticated with an API key,
// routed to the best machine that has the model, and the machine's answer is streamed straight back.

const (
	maxRequestBody   = 4 << 20
	startTimeoutWarm = 180 * time.Second // model already in memory: it should start answering quickly
	startTimeoutCold = 25 * time.Minute  // model must be loaded first (vLLM can take minutes)
	maxRecent        = 200
)

type inferJob struct {
	ch    chan Msg
	done  chan struct{}
	agent string
}

type ReqLog struct {
	TS               int64  `json:"ts"`
	Key              string `json:"key"`
	Model            string `json:"model"`
	Node             string `json:"node"`
	Status           int    `json:"status"`
	PromptTokens     int64  `json:"prompt_tokens"`
	CompletionTokens int64  `json:"completion_tokens"`
	TTFTms           int64  `json:"ttft_ms"`
	TotalMs          int64  `json:"total_ms"`
}

type recentLog struct {
	mu   sync.Mutex
	list []ReqLog
}

func (r *recentLog) add(e ReqLog) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.list = append(r.list, e)
	if len(r.list) > maxRecent {
		r.list = r.list[len(r.list)-maxRecent:]
	}
}

func (r *recentLog) snapshot() []ReqLog {
	r.mu.Lock()
	defer r.mu.Unlock()
	out := make([]ReqLog, len(r.list))
	copy(out, r.list)
	return out
}

// has reports whether a machine's set of model names contains the requested model.
// "name" also matches "name:latest" (Ollama's convention); exact names win.
func has(set map[string]bool, model string) bool {
	if set[model] {
		return true
	}
	return !strings.Contains(model, ":") && set[model+":latest"]
}

func apiError(w http.ResponseWriter, status int, msg, code string) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	json.NewEncoder(w).Encode(map[string]any{"error": map[string]any{"message": msg, "type": "wxch_error", "code": code}})
}

// pickNode chooses the machine for a model: prefer one that already has it loaded, then the least busy.
func (h *Hub) pickNode(model string) (string, *conn, bool, string) {
	h.mu.Lock()
	defer h.mu.Unlock()
	type cand struct {
		name     string
		c        *conn
		loaded   bool
		inflight int
		mem      float64
	}
	var cs []cand
	anyOnline, ownerMode := false, false
	for n, a := range h.agents {
		anyOnline = true
		b := h.stats[n]
		if b == nil || time.Since(b.at) > 60*time.Second || !has(b.inst, model) {
			continue
		}
		if b.mode != "" && b.mode != "available" { // the owner is using this machine (or paused sharing)
			ownerMode = true
			continue
		}
		cs = append(cs, cand{n, a, has(b.loaded, model), h.inflight[n], b.memAvail})
	}
	if len(cs) == 0 {
		if !anyOnline {
			return "", nil, false, "no_machine_online"
		}
		if ownerMode {
			return "", nil, false, "owner_mode"
		}
		return "", nil, false, "model_not_found"
	}
	sort.Slice(cs, func(i, j int) bool {
		if cs[i].loaded != cs[j].loaded {
			return cs[i].loaded
		}
		if cs[i].inflight != cs[j].inflight {
			return cs[i].inflight < cs[j].inflight
		}
		return cs[i].mem > cs[j].mem
	})
	return cs[0].name, cs[0].c, cs[0].loaded, ""
}

// ownerWait tells why the machines that have this model are not available and when the first one may be
// (seconds, 0 = unknown). The reason is a coarse kind (desktop, gpu, cooldown, paused), never a program name.
func (h *Hub) ownerWait(model string) (string, int) {
	h.mu.Lock()
	defer h.mu.Unlock()
	kind, wait := "", 0
	for n := range h.agents {
		b := h.stats[n]
		if b == nil || time.Since(b.at) > 60*time.Second || !has(b.inst, model) || b.mode == "" || b.mode == "available" {
			continue
		}
		w := b.waitS()
		if kind == "" || (w > 0 && (wait == 0 || w < wait)) {
			kind, wait = b.kind, w
		}
	}
	return kind, wait
}

func (h *Hub) installedModels() []string {
	h.mu.Lock()
	defer h.mu.Unlock()
	set := map[string]bool{}
	for n := range h.agents {
		if b := h.stats[n]; b != nil && time.Since(b.at) <= 60*time.Second {
			for m := range b.inst {
				set[m] = true
			}
		}
	}
	out := make([]string, 0, len(set))
	for m := range set {
		out = append(out, m)
	}
	sort.Strings(out)
	return out
}

// extractUsage finds the token counts in the tail of a response (works for JSON and SSE bodies).
// It decodes the real JSON object after the last "usage" key, so nested objects such as
// prompt_tokens_details (sent by Ollama) are handled.
func extractUsage(tail []byte) (prompt, completion int64, ok bool) {
	i := bytes.LastIndex(tail, []byte(`"usage"`))
	if i < 0 {
		return 0, 0, false
	}
	rest := bytes.TrimLeft(tail[i+len(`"usage"`):], " \t\r\n")
	if len(rest) == 0 || rest[0] != ':' {
		return 0, 0, false
	}
	var u *struct {
		P int64 `json:"prompt_tokens"`
		C int64 `json:"completion_tokens"`
	}
	if json.NewDecoder(bytes.NewReader(rest[1:])).Decode(&u) != nil || u == nil {
		return 0, 0, false
	}
	return u.P, u.C, true
}

func (s *Server) gatewayMux() http.Handler {
	mux := http.NewServeMux()
	cors := func(w http.ResponseWriter) {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Access-Control-Allow-Headers", "Authorization, Content-Type")
		w.Header().Set("Access-Control-Allow-Methods", "GET, POST, DELETE, OPTIONS")
	}
	auth := func(w http.ResponseWriter, r *http.Request) (APIKey, bool) {
		tok := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ")
		k, ok := s.keys.authenticate(tok)
		if !ok {
			apiError(w, http.StatusUnauthorized, "invalid or missing API key", "invalid_api_key")
			return k, false
		}
		return k, true
	}

	s.registerJobAPI(mux, auth, cors)

	mux.HandleFunc("/v1/models", func(w http.ResponseWriter, r *http.Request) {
		cors(w)
		if r.Method == http.MethodOptions {
			return
		}
		if _, ok := auth(w, r); !ok {
			return
		}
		type model struct {
			ID      string `json:"id"`
			Object  string `json:"object"`
			Created int64  `json:"created"`
			OwnedBy string `json:"owned_by"`
		}
		data := []model{}
		for _, m := range s.hub.installedModels() {
			data = append(data, model{m, "model", 0, "wxch"})
		}
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(map[string]any{"object": "list", "data": data})
	})

	handle := func(path string) {
		mux.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
			cors(w)
			if r.Method == http.MethodOptions {
				return
			}
			if r.Method != http.MethodPost {
				apiError(w, http.StatusMethodNotAllowed, "use POST", "method_not_allowed")
				return
			}
			s.serveInference(w, r, path)
		})
	}
	handle("/v1/chat/completions")
	handle("/v1/completions")
	handle("/v1/embeddings")

	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		cors(w)
		if r.URL.Path == "/" {
			w.Header().Set("Content-Type", "application/json")
			w.Write([]byte(`{"service":"wxch api","docs":"OpenAI-compatible: /v1/models, /v1/chat/completions, /v1/completions, /v1/embeddings, /v1/images/generations, /v1/images/edits, /v1/audio/transcriptions, /v1/audio/translations; jobs (image, video, speech): /v1/capabilities, /v1/jobs"}`))
			return
		}
		apiError(w, http.StatusNotFound, "not found", "not_found")
	})
	return mux
}

func (s *Server) serveInference(w http.ResponseWriter, r *http.Request, path string) {
	key, ok := func() (APIKey, bool) {
		tok := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ")
		k, ok := s.keys.authenticate(tok)
		if !ok {
			apiError(w, http.StatusUnauthorized, "invalid or missing API key", "invalid_api_key")
		}
		return k, ok
	}()
	if !ok {
		return
	}
	s.inferWithKey(w, r, key, path)
}

// inferWithKey is the core of a request once the caller is known (an API key, or the web chat).
func (s *Server) inferWithKey(w http.ResponseWriter, r *http.Request, key APIKey, path string) {
	doneLimit, why := s.keys.admit(key.ID)
	if doneLimit == nil {
		w.Header().Set("Retry-After", "10")
		apiError(w, http.StatusTooManyRequests, why, "rate_limit_exceeded")
		return
	}
	defer doneLimit()

	body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, maxRequestBody))
	if err != nil {
		apiError(w, http.StatusRequestEntityTooLarge, "request body too large", "request_too_large")
		return
	}
	var req map[string]any
	if json.Unmarshal(body, &req) != nil {
		apiError(w, http.StatusBadRequest, "body must be JSON", "invalid_json")
		return
	}
	model, _ := req["model"].(string)
	if model == "" {
		apiError(w, http.StatusBadRequest, "model is required", "invalid_request")
		return
	}
	stream, _ := req["stream"].(bool)
	if stream && path != "/v1/embeddings" {
		if _, has := req["stream_options"]; !has {
			req["stream_options"] = map[string]any{"include_usage": true} // so token counts come back
		}
	}
	body, _ = json.Marshal(req)

	h := s.hub
	node, a, warm, reason := h.pickNode(model)
	if a == nil {
		if reason == "no_machine_online" {
			apiError(w, http.StatusServiceUnavailable, "no machine is online", reason)
		} else if reason == "owner_mode" {
			kind, wait := h.ownerWait(model)
			if wait <= 0 {
				wait = 60 // unknown: ask the client to look again in a minute
			}
			w.Header().Set("Retry-After", strconv.Itoa(wait))
			w.Header().Set("Content-Type", "application/json")
			w.WriteHeader(http.StatusServiceUnavailable)
			json.NewEncoder(w).Encode(map[string]any{"error": map[string]any{
				"message": "the machines that have this model are being used by their owners right now; try again later",
				"type":    "wxch_error", "code": reason, "reason": kind, "retry_after_s": wait,
			}})
		} else {
			apiError(w, http.StatusNotFound, "model '"+model+"' is not available on any online machine", "model_not_found")
		}
		s.keys.record(key.ID, 0, 0, true)
		s.recent.add(ReqLog{TS: time.Now().Unix(), Key: key.Name, Model: model, Status: http.StatusNotFound})
		return
	}

	id := randHex(8)
	job := &inferJob{ch: make(chan Msg, 512), done: make(chan struct{}), agent: node}
	h.mu.Lock()
	h.infers[id] = job
	h.inflight[node]++
	h.mu.Unlock()
	defer func() {
		close(job.done)
		h.mu.Lock()
		delete(h.infers, id)
		h.inflight[node]--
		h.mu.Unlock()
	}()

	t0 := time.Now()
	a.send(Msg{T: "infer", ID: id, Method: path, Params: body})

	fl, _ := w.(http.Flusher)
	started := false
	status := 0
	var ttft int64
	var tail []byte
	events := int64(0)
	wait := startTimeoutCold
	if warm {
		wait = startTimeoutWarm
	}
	startTimer := time.NewTimer(wait)
	defer startTimer.Stop()

loop:
	for {
		select {
		case m := <-job.ch:
			switch m.T {
			case "infer_start":
				started = true
				status = m.Status
				ttft = time.Since(t0).Milliseconds()
				ct := m.Method
				if ct == "" {
					ct = "application/json"
				}
				w.Header().Set("Content-Type", ct)
				if strings.HasPrefix(ct, "text/event-stream") {
					w.Header().Set("Cache-Control", "no-cache")
					w.Header().Set("X-Accel-Buffering", "no")
				}
				w.Header().Set("X-Wxch-Node", node)
				w.WriteHeader(status)
			case "infer_chunk":
				b, err := base64.StdEncoding.DecodeString(m.D)
				if err != nil {
					continue
				}
				if ttft == 0 {
					ttft = time.Since(t0).Milliseconds()
				}
				events += int64(bytes.Count(b, []byte("data: ")))
				tail = append(tail, b...)
				if len(tail) > 16<<10 {
					tail = tail[len(tail)-(16<<10):]
				}
				w.Write(b)
				if fl != nil {
					fl.Flush()
				}
			case "infer_end":
				if !started {
					status = http.StatusBadGateway
					apiError(w, status, "machine ended the request without an answer: "+m.Err, "node_error")
				}
				break loop
			}
		case <-r.Context().Done():
			a.send(Msg{T: "infer_cancel", ID: id})
			status = 499
			break loop
		case <-startTimer.C:
			if !started {
				a.send(Msg{T: "infer_cancel", ID: id})
				status = http.StatusGatewayTimeout
				apiError(w, status, "machine did not answer in time", "node_timeout")
				break loop
			}
		}
	}

	pt, ct, found := extractUsage(tail)
	if !found && events > 1 {
		ct = events - 1 // streaming without usage: roughly one event per token, minus the final [DONE]
	}
	failed := status >= 400
	s.keys.record(key.ID, pt, ct, failed)
	s.recent.add(ReqLog{TS: time.Now().Unix(), Key: key.Name, Model: model, Node: node, Status: status,
		PromptTokens: pt, CompletionTokens: ct, TTFTms: ttft, TotalMs: time.Since(t0).Milliseconds()})
}

// ---- management API for the web UI (cookie login) ----

func (s *Server) keysAPI(w http.ResponseWriter, r *http.Request) {
	if !s.authed(r) {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	rest := strings.Trim(strings.TrimPrefix(r.URL.Path, "/api/keys"), "/")
	switch {
	case rest == "" && r.Method == http.MethodGet:
		json.NewEncoder(w).Encode(s.keys.list())
	case rest == "" && r.Method == http.MethodPost:
		var req struct{ Name, Scope string }
		if json.NewDecoder(http.MaxBytesReader(w, r.Body, 1024)).Decode(&req) != nil || strings.TrimSpace(req.Name) == "" || len(req.Name) > 60 || (req.Scope != "" && req.Scope != "talk") {
			http.Error(w, "name required (max 60 chars)", http.StatusBadRequest)
			return
		}
		key, k := s.keys.createScoped(strings.TrimSpace(req.Name), req.Scope)
		json.NewEncoder(w).Encode(map[string]any{"key": key, "id": k.ID, "name": k.Name})
	case rest == "recent" && r.Method == http.MethodGet:
		json.NewEncoder(w).Encode(s.recent.snapshot())
	case strings.HasSuffix(rest, "/disable") && r.Method == http.MethodPost:
		id := strings.TrimSuffix(rest, "/disable")
		if !s.keys.setDisabled(id, r.URL.Query().Get("on") == "1") {
			http.NotFound(w, r)
			return
		}
		w.Write([]byte(`"ok"`))
	case r.Method == http.MethodDelete:
		if !s.keys.remove(rest) {
			http.NotFound(w, r)
			return
		}
		w.Write([]byte(`"ok"`))
	default:
		http.Error(w, "bad request", http.StatusBadRequest)
	}
}