package main

import (
	"encoding/json"
	"fmt"
	"io"
	"io/fs"
	"net/http"
	"os"
	"path"
	"regexp"
	"sort"
	"strings"
)

// pageHandler makes browsers re-check the embedded pages on every load, so a new
// deployment is picked up immediately instead of serving a stale copy.
func pageHandler(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Cache-Control", "no-cache")
		next.ServeHTTP(w, r)
	})
}

// baseURL is the public address of this server as the caller reached it
// (or the configured public_url). Installer scripts use it instead of a fixed domain.
func baseURL(r *http.Request, configured string) string {
	if configured != "" {
		return strings.TrimRight(configured, "/")
	}
	scheme := "https"
	if p := r.Header.Get("X-Forwarded-Proto"); p != "" {
		scheme = p
	} else if r.TLS == nil {
		scheme = "http"
	}
	return scheme + "://" + r.Host
}

// baseOrPlaceholder is for messages printed by the command line, where there is no request.
func baseOrPlaceholder(configured string) string {
	if configured == "" {
		return "https://<your-control-host>"
	}
	return strings.TrimRight(configured, "/")
}

// agentTokens reads the machines from the config file, so machines added while running are seen
// (same as wsAgent).
func (s *Server) agentTokens() map[string]string {
	tokens := s.cfg.AgentTokens
	if b, err := os.ReadFile(s.cfgPath); err == nil {
		var c Config
		if json.Unmarshal(b, &c) == nil {
			tokens = c.AgentTokens
		}
	}
	return tokens
}

// agentsRegistered lists the names of the machines that have a token, online or not (never the tokens),
// so the page can show a machine that has not connected yet.
func (s *Server) agentsRegistered(w http.ResponseWriter, r *http.Request) {
	if !s.authed(r) {
		http.Error(w, "no", 401)
		return
	}
	names := []string{}
	for n := range s.agentTokens() {
		names = append(names, n)
	}
	sort.Strings(names)
	w.Header().Set("Cache-Control", "no-store")
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(names)
}

// agentCmd returns the one-line install command for an existing machine, so the owner can
// copy it from the UI. It carries the machine's token, so it is login-only and never cached.
func (s *Server) agentCmd(w http.ResponseWriter, r *http.Request) {
	if !s.authed(r) {
		http.Error(w, "no", 401)
		return
	}
	tok, ok := s.agentTokens()[r.URL.Query().Get("name")]
	if !ok {
		http.NotFound(w, r)
		return
	}
	w.Header().Set("Cache-Control", "no-store")
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	fmt.Fprintf(w, "curl -fsSL %s/get | WXCH_TOKEN=%s sh", baseURL(r, s.cfg.PublicURL), tok)
}

var stepNameRe = regexp.MustCompile(`^[a-z0-9-]{1,40}\.sh$`)

// registerInstallers serves the installer scripts. @@BASE@@ inside a script is replaced
// with this server's own address, so the same scripts work on any deployment.
//
//	/get          one-line bootstrap: curl -fsSL <base>/get | sh   (asks for sudo once)
//	/setup.sh     the installer itself (runs as root)
//	/uninstall.sh removes what the installer added
//	/add-disk.sh  adds an external model disk (run through /get ... add-disk)
//	/install.sh   agent-only installer (used to upgrade an existing machine)
//	/steps/X.sh   the individual steps setup.sh runs
//	/lab/X.sh     scripts for trying things by hand, such as the ComfyUI experiment
func registerInstallers(mux *http.ServeMux, files fs.FS, configured string) {
	serve := func(name string) http.HandlerFunc {
		return func(w http.ResponseWriter, r *http.Request) {
			b, err := fs.ReadFile(files, name)
			if err != nil {
				http.NotFound(w, r)
				return
			}
			out := strings.ReplaceAll(string(b), "@@BASE@@", baseURL(r, configured))
			w.Header().Set("Content-Type", "text/x-shellscript; charset=utf-8")
			w.Header().Set("Cache-Control", "no-cache")
			io.WriteString(w, out)
		}
	}
	mux.HandleFunc("/get", serve("get.sh"))
	mux.HandleFunc("/setup.sh", serve("setup.sh"))
	mux.HandleFunc("/uninstall.sh", serve("uninstall.sh"))
	mux.HandleFunc("/add-disk.sh", serve("add-disk.sh"))
	mux.HandleFunc("/install.sh", serve("install-agent.sh"))
	// scripts for trying things by hand (not part of the installer); the file name is checked like a step's
	mux.HandleFunc("/lab/", func(w http.ResponseWriter, r *http.Request) {
		name := path.Base(r.URL.Path)
		if !stepNameRe.MatchString(name) {
			http.NotFound(w, r)
			return
		}
		serve("lab/"+name)(w, r)
	})
	mux.HandleFunc("/steps/", func(w http.ResponseWriter, r *http.Request) {
		name := path.Base(r.URL.Path)
		if !stepNameRe.MatchString(name) {
			http.NotFound(w, r)
			return
		}
		serve("steps/"+name)(w, r)
	})
}