package main

import (
	"crypto/subtle"
	"embed"
	"encoding/json"
	"fmt"
	"io/fs"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

//go:embed webui/* webui_glass/*
var webAssets embed.FS

// webThemes maps a user-facing theme name to its embedded asset directory. Both
// designs share the same REST API and JS logic; only the static shell differs,
// so a theme is just a different sub-directory of the embedded FS.
var webThemes = map[string]string{
	"classic": "webui",       // original Modern Dark design (default)
	"glass":   "webui_glass", // glassmorphism · Slate + Amber
}

// webModes are the accepted colour modes. "dark" (default) forces the dark
// palette; "light" forces light; "auto" follows the viewer's OS preference.
// The mode is stamped onto <html data-theme="…"> at serve time — see
// injectUIMode. glass is a dark-only design and ignores light/auto.
var webModes = map[string]bool{"dark": true, "light": true, "auto": true}

// injectUIMode stamps data-theme="<mode>" onto the <html> tag of the served
// index.html so the stylesheet resolves the palette before first paint. It
// keys off the `<html lang="en"` prefix both theme shells share; if that
// anchor is ever renamed the page still works (it just falls back to CSS
// defaults), so the replace is best-effort.
func injectUIMode(html []byte, mode string) []byte {
	const anchor = `<html lang="en"`
	attr := fmt.Sprintf(`<html lang="en" data-theme=%q`, mode)
	return []byte(strings.Replace(string(html), anchor, attr, 1))
}

// opLock serialises all mutating operations (create/delete/enable/disable/sync
// /server control). The underlying registry + interface state are shared, so we
// process one change at a time to avoid races between concurrent requests.
var opLock sync.Mutex

// runWeb starts the management web server. It is deliberately the only place
// that flips webMode on, so every die() reached from a handler becomes a
// recoverable panic instead of killing the process.
func runWeb(args []string) {
	addr := "127.0.0.1:8080"
	// Theme precedence: --theme flag > AWG_WEB_THEME env > "classic" default.
	theme := os.Getenv("AWG_WEB_THEME")
	if theme == "" {
		theme = "classic"
	}
	// Colour-mode precedence: --ui-mode flag > AWG_WEB_MODE env > "dark" default.
	// "dark" forces the dark palette regardless of the viewer's OS preference.
	uiMode := os.Getenv("AWG_WEB_MODE")
	if uiMode == "" {
		uiMode = "dark"
	}
	for i := 0; i < len(args); i++ {
		switch args[i] {
		case "-addr", "--addr":
			if i+1 >= len(args) {
				die("--addr requires a value (e.g. 0.0.0.0:8080)")
			}
			addr = args[i+1]
			i++
		case "-theme", "--theme":
			if i+1 >= len(args) {
				die("--theme requires a value (classic or glass)")
			}
			theme = args[i+1]
			i++
		case "-ui-mode", "--ui-mode":
			if i+1 >= len(args) {
				die("--ui-mode requires a value (dark, light or auto)")
			}
			uiMode = args[i+1]
			i++
		default:
			die("Unknown web option: %s", args[i])
		}
	}

	assetDir, ok := webThemes[theme]
	if !ok {
		die("Unknown theme %q (choose: classic, glass)", theme)
	}
	if !webModes[uiMode] {
		die("Unknown ui-mode %q (choose: dark, light, auto)", uiMode)
	}

	webMode = true

	// Accumulate traffic totals in the background so counters keep folding into
	// the durable store even when no browser is polling.
	go statsSampler()

	user := os.Getenv("AWG_WEB_USER")
	pass := os.Getenv("AWG_WEB_PASS")
	authOn := user != "" && pass != ""

	mux := http.NewServeMux()
	registerAPI(mux)

	// Static SPA assets, served from the selected theme's embedded directory.
	sub, err := fs.Sub(webAssets, assetDir)
	if err != nil {
		die("failed to open embedded assets: %v", err)
	}
	// index.html is served with the colour mode stamped onto <html> so the CSS
	// applies it before first paint (no flash, no client JS). All other assets
	// go through the plain file server.
	rawIndex, err := fs.ReadFile(sub, "index.html")
	if err != nil {
		die("failed to read embedded index.html: %v", err)
	}
	indexHTML := injectUIMode(rawIndex, uiMode)
	fileSrv := http.FileServer(http.FS(sub))
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/" || r.URL.Path == "/index.html" {
			w.Header().Set("Content-Type", "text/html; charset=utf-8")
			_, _ = w.Write(indexHTML)
			return
		}
		fileSrv.ServeHTTP(w, r)
	})

	handler := basicAuth(authOn, user, pass, mux)

	srv := &http.Server{
		Addr:              addr,
		Handler:           handler,
		ReadHeaderTimeout: 15 * time.Second,
		// No write timeout: install-deps / init-server can run for minutes.
	}

	fmt.Printf("INFO: AmneziaWG web UI listening on http://%s (theme: %s, ui-mode: %s)\n", addr, theme, uiMode)
	if authOn {
		fmt.Println("INFO: HTTP Basic auth enabled (AWG_WEB_USER/AWG_WEB_PASS)")
	} else {
		fmt.Println("WARN: no auth set — bind to 127.0.0.1 or set AWG_WEB_USER/AWG_WEB_PASS")
	}
	if err := srv.ListenAndServe(); err != nil {
		die("web server error: %v", err)
	}
}

// statsSampler periodically folds live counters into the persisted totals so
// accumulation continues with no browser open and counter resets are caught
// promptly. Each tick is isolated: a panic (e.g. a die() from loadConfig in web
// mode) is recovered so the background loop can never crash the server.
func statsSampler() {
	for range time.Tick(20 * time.Second) {
		func() {
			defer func() { _ = recover() }()
			if !serverInitialized() {
				return
			}
			c := loadConfig()
			sampleStats(peerStats(c.Interface))
		}()
	}
}

// basicAuth optionally guards every request with HTTP Basic credentials.
func basicAuth(on bool, user, pass string, next http.Handler) http.Handler {
	if !on {
		return next
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		u, p, ok := r.BasicAuth()
		userOK := subtle.ConstantTimeCompare([]byte(u), []byte(user)) == 1
		passOK := subtle.ConstantTimeCompare([]byte(p), []byte(pass)) == 1
		if !ok || !userOK || !passOK {
			w.Header().Set("WWW-Authenticate", `Basic realm="awg-profiler"`)
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}
		next.ServeHTTP(w, r)
	})
}

// ─── routing ─────────────────────────────────────────────────────────────────────

func registerAPI(mux *http.ServeMux) {
	mux.HandleFunc("GET /api/setup-status", h(handleSetupStatus))
	mux.HandleFunc("GET /api/status", h(handleStatus))
	mux.HandleFunc("GET /api/clients", h(handleClients))
	mux.HandleFunc("POST /api/clients", h(handleCreate))
	mux.HandleFunc("DELETE /api/clients/{id}", h(handleDelete))
	mux.HandleFunc("POST /api/clients/{id}/enable", h(handleEnable))
	mux.HandleFunc("POST /api/clients/{id}/disable", h(handleDisable))
	mux.HandleFunc("POST /api/clients/{id}/comment", h(handleComment))
	mux.HandleFunc("POST /api/clients/{id}/stats/reset", h(handleStatsReset))
	mux.HandleFunc("GET /api/clients/{id}/config", h(handleConfigDownload))
	mux.HandleFunc("GET /api/clients/{id}/qr", h(handleQR))
	mux.HandleFunc("POST /api/server/start", h(handleServerStart))
	mux.HandleFunc("POST /api/server/stop", h(handleServerStop))
	mux.HandleFunc("POST /api/server/restart", h(handleServerRestart))
	mux.HandleFunc("POST /api/server/sync", h(handleSync))
	mux.HandleFunc("POST /api/install-deps", h(handleInstallDeps))
	mux.HandleFunc("POST /api/init-server", h(handleInitServer))
}

// h wraps an API handler with CSRF hardening and panic recovery (turning
// die() into a JSON error response).
func h(fn func(w http.ResponseWriter, r *http.Request)) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if !csrfSafe(r) {
			writeErr(w, http.StatusUnsupportedMediaType, "Content-Type must be application/json")
			return
		}
		defer func() {
			if rec := recover(); rec != nil {
				msg := fmt.Sprintf("%v", rec)
				if de, ok := rec.(dieError); ok {
					msg = de.msg
				}
				writeErr(w, http.StatusBadRequest, msg)
			}
		}()
		fn(w, r)
	}
}

// csrfSafe blocks cross-site form-triggered state changes. Basic-auth
// credentials are attached to same-origin requests automatically by the
// browser, so without this check a malicious page could submit a blind
// <form> POST/DELETE (forms can only send text/plain,
// application/x-www-form-urlencoded, or multipart/form-data — never
// application/json) and trigger e.g. server-stop or client deletion under
// the logged-in admin's session. GET requests are read-only and exempt.
func csrfSafe(r *http.Request) bool {
	if r.Method != http.MethodPost && r.Method != http.MethodDelete {
		return true
	}
	return strings.HasPrefix(r.Header.Get("Content-Type"), "application/json")
}

// ─── response helpers ──────────────────────────────────────────────────────────

func writeJSON(w http.ResponseWriter, code int, v any) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.WriteHeader(code)
	_ = json.NewEncoder(w).Encode(v)
}

func writeErr(w http.ResponseWriter, code int, msg string) {
	writeJSON(w, code, map[string]string{"error": msg})
}

// serverInitialized reports whether init-server has produced a config, without
// dying (used to decide between the dashboard and the setup wizard).
func serverInitialized() bool {
	_, err := os.Stat(configFile)
	return err == nil
}

// ─── handlers: status / setup ───────────────────────────────────────────────────

func handleSetupStatus(w http.ResponseWriter, r *http.Request) {
	st := loadState()
	resp := map[string]any{
		"deps_installed":     st.DepsInstalled,
		"deps_installed_at":  st.InstalledAt,
		"server_initialized": serverInitialized(),
	}
	writeJSON(w, http.StatusOK, resp)
}

func handleStatus(w http.ResponseWriter, r *http.Request) {
	if !serverInitialized() {
		writeJSON(w, http.StatusOK, map[string]any{"initialized": false})
		return
	}
	c := loadConfig()
	up := awgIfaceUp(c.Interface)

	clients := loadRegistry()
	active := 0
	for _, cl := range clients {
		if cl.IsEnabled == "ACTIVE" {
			active++
		}
	}
	stats := peerStats(c.Interface)
	sampleStats(stats) // keep totals accumulating even on status-only polls
	online := 0
	for _, s := range stats {
		if s.Online() {
			online++
		}
	}

	writeJSON(w, http.StatusOK, map[string]any{
		"initialized":    true,
		"interface":      c.Interface,
		"network":        c.Network,
		"port":           c.Port,
		"public_ip":      c.PublicIP,
		"public_key":     c.ServerPub,
		"dns":            c.DNS,
		"mtu":            c.MTU,
		"interface_up":   up,
		"total_clients":  len(clients),
		"active_clients": active,
		"online_clients": online,
	})
}

// ─── handlers: clients ──────────────────────────────────────────────────────────

// clientOut is the browser-facing client shape. It deliberately omits the
// private and preshared keys — those only ever leave via the .conf download.
type clientOut struct {
	ID              int64  `json:"id"`
	Name            string `json:"name"`
	IP              string `json:"ip"`
	PublicKey       string `json:"public_key"`
	IsEnabled       string `json:"is_enabled"`
	CreatedAt       int64  `json:"created_at"`
	Comment         string `json:"comment"`
	Online          bool   `json:"online"`
	Endpoint        string `json:"endpoint"`
	LatestHandshake int64  `json:"latest_handshake"`
	TransferRx      int64  `json:"transfer_rx"`
	TransferTx      int64  `json:"transfer_tx"`
	// Cumulative, restart-surviving traffic accumulated from the stats store,
	// plus the mark from which it has been counting. These are what the UI shows.
	TotalRx    int64 `json:"total_rx"`
	TotalTx    int64 `json:"total_tx"`
	StatsSince int64 `json:"stats_since"`
}

func handleClients(w http.ResponseWriter, r *http.Request) {
	c := loadConfig()
	initStorage()
	clients := loadRegistry()
	stats := peerStats(c.Interface)
	// Fold this live reading into the persisted totals and get the fresh snapshot.
	totals := sampleStats(stats)

	out := make([]clientOut, 0, len(clients))
	for _, cl := range clients {
		o := clientOut{
			ID:        cl.ID,
			Name:      cl.Name,
			IP:        cl.IP,
			PublicKey: cl.PublicKey,
			IsEnabled: cl.IsEnabled,
			CreatedAt: cl.CreatedAt,
			Comment:   cl.Comment,
		}
		if s, ok := stats[cl.PublicKey]; ok {
			o.Online = s.Online()
			o.Endpoint = s.Endpoint
			o.LatestHandshake = s.LatestHandshake
			o.TransferRx = s.TransferRx
			o.TransferTx = s.TransferTx
		}
		if rec, ok := totals[cl.PublicKey]; ok {
			o.TotalRx = rec.TotalRx
			o.TotalTx = rec.TotalTx
			o.StatsSince = rec.Since
		}
		out = append(out, o)
	}
	writeJSON(w, http.StatusOK, out)
}

func handleCreate(w http.ResponseWriter, r *http.Request) {
	var body struct {
		Name string `json:"name"`
	}
	if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
		writeErr(w, http.StatusBadRequest, "invalid JSON body")
		return
	}
	name := strings.TrimSpace(body.Name)
	if name == "" {
		writeErr(w, http.StatusBadRequest, "client name required")
		return
	}

	opLock.Lock()
	defer opLock.Unlock()

	requireBinary("awg")
	requireBinary("qrencode")
	c := loadConfig()
	initStorage()
	createClient(c, detectOS(), []string{name})

	// Return the freshly created client (highest id with this sanitized name).
	clients := loadRegistry()
	sanitized := sanitize(name)
	var created *Client
	for i := range clients {
		if clients[i].Name == sanitized {
			created = &clients[i]
		}
	}
	if created == nil {
		writeErr(w, http.StatusInternalServerError, "client created but not found in registry")
		return
	}
	writeJSON(w, http.StatusCreated, clientOut{
		ID: created.ID, Name: created.Name, IP: created.IP,
		PublicKey: created.PublicKey, IsEnabled: created.IsEnabled,
		CreatedAt: created.CreatedAt, Comment: created.Comment,
	})
}

func handleComment(w http.ResponseWriter, r *http.Request) {
	var body struct {
		Comment string `json:"comment"`
	}
	if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
		writeErr(w, http.StatusBadRequest, "invalid JSON body")
		return
	}
	opLock.Lock()
	defer opLock.Unlock()
	loadConfig()
	initStorage()
	id := parseID([]string{r.PathValue("id")})
	cl := setClientComment(id, strings.TrimSpace(body.Comment))
	writeJSON(w, http.StatusOK, map[string]any{"status": "updated", "comment": cl.Comment})
}

// handleStatsReset clears a client's accumulated traffic and rebaselines it to
// the current live counter, so the tally restarts from zero as of now.
func handleStatsReset(w http.ResponseWriter, r *http.Request) {
	c := loadConfig()
	cl := clientByID(r.PathValue("id")) // dies (→ 400/404) if the id is unknown
	// Baseline to the current raw counter so we don't re-add pre-reset bytes. If
	// the peer is absent from the dump the interface is down (counters restart
	// near zero on the way back up), so a zero baseline is correct.
	live := peerStats(c.Interface)[cl.PublicKey]
	rec := resetStats(cl.PublicKey, live.TransferRx, live.TransferTx)
	writeJSON(w, http.StatusOK, map[string]any{
		"status":      "reset",
		"total_rx":    rec.TotalRx,
		"total_tx":    rec.TotalTx,
		"stats_since": rec.Since,
	})
}

func handleDelete(w http.ResponseWriter, r *http.Request) {
	opLock.Lock()
	defer opLock.Unlock()
	c := loadConfig()
	initStorage()
	deleteClient(c, detectOS(), []string{r.PathValue("id")})
	writeJSON(w, http.StatusOK, map[string]string{"status": "deleted"})
}

func handleEnable(w http.ResponseWriter, r *http.Request) {
	opLock.Lock()
	defer opLock.Unlock()
	c := loadConfig()
	initStorage()
	enableClient(c, detectOS(), []string{r.PathValue("id")})
	writeJSON(w, http.StatusOK, map[string]string{"status": "enabled"})
}

func handleDisable(w http.ResponseWriter, r *http.Request) {
	opLock.Lock()
	defer opLock.Unlock()
	c := loadConfig()
	initStorage()
	disableClient(c, detectOS(), []string{r.PathValue("id")})
	writeJSON(w, http.StatusOK, map[string]string{"status": "disabled"})
}

// clientByID looks up a client for download endpoints.
func clientByID(idStr string) *Client {
	id := parseID([]string{idStr})
	clients := loadRegistry()
	cl := findClient(clients, id)
	if cl == nil {
		die("Client with ID %d not found", id)
	}
	return cl
}

func handleConfigDownload(w http.ResponseWriter, r *http.Request) {
	cl := clientByID(r.PathValue("id"))
	path := filepath.Join(clientDir, cl.Name+".conf")
	data, err := os.ReadFile(path)
	if err != nil {
		die("config file for %s not found", cl.Name)
	}
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	w.Header().Set("Content-Disposition", fmt.Sprintf(`attachment; filename="%s.conf"`, cl.Name))
	_, _ = w.Write(data)
}

func handleQR(w http.ResponseWriter, r *http.Request) {
	cl := clientByID(r.PathValue("id"))
	path := filepath.Join(clientDir, cl.Name+".png")
	data, err := os.ReadFile(path)
	if err != nil {
		die("QR image for %s not found", cl.Name)
	}
	w.Header().Set("Content-Type", "image/png")
	_, _ = w.Write(data)
}

// ─── handlers: server control ───────────────────────────────────────────────────

func handleServerStart(w http.ResponseWriter, r *http.Request) {
	opLock.Lock()
	defer opLock.Unlock()
	c := loadConfig()
	serviceStart(detectOS(), c.Interface)
	writeJSON(w, http.StatusOK, map[string]string{"status": "started"})
}

func handleServerStop(w http.ResponseWriter, r *http.Request) {
	opLock.Lock()
	defer opLock.Unlock()
	c := loadConfig()
	serviceStop(detectOS(), c.Interface)
	writeJSON(w, http.StatusOK, map[string]string{"status": "stopped"})
}

func handleServerRestart(w http.ResponseWriter, r *http.Request) {
	opLock.Lock()
	defer opLock.Unlock()
	c := loadConfig()
	serviceRestart(detectOS(), c.Interface)
	writeJSON(w, http.StatusOK, map[string]string{"status": "restarted"})
}

func handleSync(w http.ResponseWriter, r *http.Request) {
	var body struct {
		Restart bool `json:"restart"`
	}
	_ = json.NewDecoder(r.Body).Decode(&body)

	opLock.Lock()
	defer opLock.Unlock()
	c := loadConfig()
	syncConfigWeb(c, detectOS(), body.Restart)
	writeJSON(w, http.StatusOK, map[string]any{"status": "synced", "restarted": body.Restart})
}

// ─── handlers: setup wizard ─────────────────────────────────────────────────────

func handleInstallDeps(w http.ResponseWriter, r *http.Request) {
	opLock.Lock()
	defer opLock.Unlock()
	installDeps() // logs stream to the server console; panics (recovered) on failure
	writeJSON(w, http.StatusOK, map[string]string{"status": "installed"})
}

func handleInitServer(w http.ResponseWriter, r *http.Request) {
	var p initParams
	if err := json.NewDecoder(r.Body).Decode(&p); err != nil {
		writeErr(w, http.StatusBadRequest, "invalid JSON body")
		return
	}
	opLock.Lock()
	defer opLock.Unlock()
	if serverInitialized() {
		writeErr(w, http.StatusConflict, "server already initialized")
		return
	}
	c := initServerWeb(p)
	writeJSON(w, http.StatusOK, map[string]any{
		"status":     "initialized",
		"interface":  c.Interface,
		"public_key": c.ServerPub,
	})
}
