package main

import (
	"encoding/json"
	"io/fs"
	"log"
	"net/http"
	"sort"
	"time"

	"golang.org/x/crypto/bcrypt"
)

// Web chat on chat.wxch.org. Login is the web password (like the notepad); the page talks to
// /chat/completions, which goes through the same gateway core as API keys, under the name "web-chat".

var webChatKey = APIKey{ID: "webchat", Name: "web-chat"}

func (s *Server) chatAuthed(r *http.Request) bool {
	c, err := r.Cookie("wxch_chat")
	if err != nil {
		return false
	}
	return s.chatSess.valid(c.Value)
}

type chatModel struct {
	Name      string `json:"name"`
	Size      int64  `json:"size"`
	Loaded    bool   `json:"loaded"`    // in memory on at least one machine (fast first reply)
	Available bool   `json:"available"` // at least one machine can take it right now
	Nodes     int    `json:"nodes"`
	Runtime   string `json:"runtime"`           // ollama | vllm | lmstudio
	Kind      string `json:"kind,omitempty"`    // not available: desktop | gpu | cooldown | paused
	RetryS    int    `json:"retry_s,omitempty"` // not available: estimated seconds until the first machine is free (0 = unknown)
}

// chatModels lists what can be chatted with, from the latest telemetry of online machines.
func (h *Hub) chatModels() []chatModel {
	h.mu.Lock()
	defer h.mu.Unlock()
	byName := map[string]*chatModel{}
	for n := range h.agents {
		b := h.stats[n]
		if b == nil || time.Since(b.at) > 60*time.Second {
			continue
		}
		var p struct {
			Installed []struct {
				Name    string `json:"name"`
				Size    int64  `json:"size"`
				Runtime string `json:"runtime"`
			} `json:"installed"`
		}
		json.Unmarshal(b.latest, &p)
		for _, m := range p.Installed {
			name := m.Name
			cm := byName[name]
			if cm == nil {
				cm = &chatModel{Name: name, Size: m.Size, Runtime: m.Runtime}
				byName[name] = cm
			}
			cm.Nodes++
			if b.loaded[name] {
				cm.Loaded = true
			}
			if b.mode == "" || b.mode == "available" {
				cm.Available = true
			} else if w := b.waitS(); cm.Kind == "" || (w > 0 && (cm.RetryS == 0 || w < cm.RetryS)) {
				cm.Kind, cm.RetryS = b.kind, w // the machine that frees up first
			}
		}
	}
	out := make([]chatModel, 0, len(byName))
	for _, m := range byName {
		if m.Available {
			m.Kind, m.RetryS = "", 0
		}
		out = append(out, *m)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Name < out[j].Name })
	return out
}

func (s *Server) chatMux(sub fs.FS) http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/" {
			http.NotFound(w, r)
			return
		}
		b, _ := fs.ReadFile(sub, "chat.html")
		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		w.Header().Set("Cache-Control", "no-store")
		w.Write(b)
	})
	mux.HandleFunc("/chat/login", func(w http.ResponseWriter, r *http.Request) {
		ip := clientIP(r)
		if s.locked(ip) {
			http.Error(w, "too many attempts", http.StatusTooManyRequests)
			return
		}
		var req struct{ Password string }
		if json.NewDecoder(http.MaxBytesReader(w, r.Body, 4096)).Decode(&req) != nil ||
			bcrypt.CompareHashAndPassword([]byte(s.pwHash()), []byte(req.Password)) != nil {
			s.fail(ip)
			log.Printf("chat login failed from %s", ip)
			http.Error(w, "invalid", http.StatusUnauthorized)
			return
		}
		tok := randHex(32)
		s.chatSess.add(tok, 30*24*time.Hour)
		http.SetCookie(w, &http.Cookie{Name: "wxch_chat", Value: tok, Path: "/", HttpOnly: true,
			Secure: true, SameSite: http.SameSiteStrictMode, MaxAge: 30 * 24 * 3600})
		w.Write([]byte("ok"))
	})
	mux.HandleFunc("/chat/models", func(w http.ResponseWriter, r *http.Request) {
		if !s.chatAuthed(r) {
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		w.Header().Set("Cache-Control", "no-store")
		json.NewEncoder(w).Encode(s.hub.chatModels())
	})
	mux.HandleFunc("/chat/completions", func(w http.ResponseWriter, r *http.Request) {
		if !s.chatAuthed(r) {
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		if r.Method != http.MethodPost {
			http.Error(w, "use POST", http.StatusMethodNotAllowed)
			return
		}
		s.inferWithKey(w, r, webChatKey, "/v1/chat/completions")
	})
	return mux
}