//go:build !(js && wasm)

// Package wasmdevserver is a reusable development server for gowasm apps: it serves
// the built web assets, renders routes server-side (SSR) at request time, hosts
// the /rsc server-component endpoint, and hot-swaps the freshly built wasm into
// the browser on change (no full reload, state preserved) — surfacing build
// failures as an in-page overlay.
//
// It never imports application code (per kjol's golden rule). The app injects
// everything specific to it through Config: how to build the wasm bundle
// (Build), how to render a route to HTML (Render), and how to wrap that HTML in
// a document (Document). The WebSocket hub and file watcher use only the
// standard library.
package wasmdevserver

import (
	"crypto/sha1"
	"encoding/base64"
	"encoding/json"
	"io"
	"io/fs"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"kjol/wasmruntime/rsc"
)

// Config wires an app into the dev server. Render and Document are called per
// request; Build is called for the initial build and on every source change.
type Config struct {
	Addr      string                 // listen address (default ":8085")
	Dir       string                 // static assets dir; Build writes app.wasm here (default "./wwwroot")
	Watch     bool                   // rebuild on change + hot reload
	WatchDirs []string               // source dirs to watch when Watch is set
	Build     func() ([]byte, error) // (re)build everything; combined output on failure

	// BuildCSS recompiles ONLY the stylesheet. When a save touched stylesheets and no
	// Go code, this runs instead of Build and the browser swaps the stylesheet in
	// place — no wasm rebuild, no page reload, no lost state. Without it a colour
	// tweak waits on the whole Go compiler, which is the difference between a design
	// loop and a coffee break.
	BuildCSS func() ([]byte, error)
	Render   func(path string) (inner string, ok bool) // SSR the #app inner HTML for a route (ok=false => client-rendered)
	Document func(inner string) string                 // wrap #app inner HTML in a full HTML document
	Handle   func(mux *http.ServeMux)                  // optional: register extra routes (e.g. app API endpoints)
}

// Serve builds once (in watch mode), wires the routes, and blocks serving.
func Serve(cfg Config) error {
	if cfg.Addr == "" {
		cfg.Addr = ":8085"
	}
	if cfg.Dir == "" {
		cfg.Dir = "./wwwroot"
	}
	mux := http.NewServeMux()
	h := newHub()

	if cfg.Watch {
		if err := ensureShim(cfg.Dir); err != nil {
			log.Printf("warning: could not stage wasm_exec.js: %v", err)
		}
		if cfg.Build != nil {
			if out, err := cfg.Build(); err != nil {
				log.Printf("initial build failed: %v\n%s", err, out)
				h.setError(string(out)) // a browser opened now sees it via the overlay
			}
		}
		mux.HandleFunc("/livereload", h.serveWS)
		mux.HandleFunc("/livereload.js", serveClientJS)
		go watchLoop(cfg, h)
		log.Printf("hot reload enabled (watching %v)", cfg.WatchDirs)
	}

	mux.HandleFunc("POST /rsc", rsc.Handler) // server components
	if cfg.Handle != nil {
		cfg.Handle(mux) // app-registered routes (API endpoints, etc.)
	}
	mux.HandleFunc("/", rootHandler(cfg))

	log.Printf("serving %q on http://localhost%s", cfg.Dir, cfg.Addr)
	return http.ListenAndServe(cfg.Addr, mux)
}

// ---- serving: assets, and dynamic SSR for HTML routes -------------------

func rootHandler(cfg Config) http.HandlerFunc {
	live := ""
	if cfg.Watch {
		live = `<script src="/livereload.js"></script>`
	}
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Cache-Control", "no-store")
		clean := filepath.Clean("/" + r.URL.Path)
		fsPath := filepath.Join(cfg.Dir, clean)
		if info, err := os.Stat(fsPath); err == nil && !info.IsDir() && clean != "/" {
			if strings.HasSuffix(fsPath, ".wasm") {
				w.Header().Set("Content-Type", "application/wasm")
			}
			http.ServeFile(w, r, fsPath)
			return
		}
		inner := ""
		if cfg.Render != nil {
			inner, _ = cfg.Render(r.URL.Path)
		}
		html := inner
		if cfg.Document != nil {
			html = cfg.Document(inner)
		}
		if live != "" {
			html = strings.Replace(html, "</body>", live+"\n</body>", 1)
		}
		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		io.WriteString(w, html)
	}
}

// ---- build + watch ------------------------------------------------------

// watchLoop rebuilds on change, and treats a stylesheet edit differently from a code
// edit — because they cost wildly different amounts.
//
// A Go change means recompiling a multi-megabyte wasm binary. A CSS change means
// re-running Tailwind, which takes a moment. Putting both through the same path
// would make every tweak to a colour wait on the compiler, so a save that touched
// ONLY stylesheets runs BuildCSS and pushes a stylesheet swap: no wasm rebuild, no
// page reload, no lost state — the new CSS just appears.
func watchLoop(cfg Config, h *hub) {
	prevGo := fingerprintExt(cfg.WatchDirs, ".go")
	prevCSS := fingerprintExt(cfg.WatchDirs, ".css")

	for {
		time.Sleep(300 * time.Millisecond)

		fpGo := fingerprintExt(cfg.WatchDirs, ".go")
		fpCSS := fingerprintExt(cfg.WatchDirs, ".css")
		codeChanged := fpGo != prevGo
		cssChanged := fpCSS != prevCSS
		if !codeChanged && !cssChanged {
			continue
		}
		prevGo, prevCSS = fpGo, fpCSS

		// Styles only: the cheap path.
		if !codeChanged && cfg.BuildCSS != nil {
			log.Println("stylesheet changed, recompiling CSS…")
			h.broadcast(`{"type":"building"}`)
			if out, err := cfg.BuildCSS(); err != nil {
				log.Printf("CSS build failed: %v\n%s", err, out)
				h.setError(string(out))
				continue
			}
			log.Println("CSS ok — swapping stylesheets")
			h.clearError()
			h.broadcast(`{"type":"css"}`)
			continue
		}

		log.Println("change detected, rebuilding…")
		h.broadcast(`{"type":"building"}`)
		if cfg.Build == nil {
			continue
		}
		if out, err := cfg.Build(); err != nil {
			log.Printf("build failed: %v\n%s", err, out)
			h.setError(string(out)) // push the compiler output to the browser overlay
			continue
		}
		log.Println("rebuild ok — reloading clients")
		h.clearError()
		h.broadcast(`{"type":"reload"}`)
	}
}

// fingerprintExt changes whenever a file of the given extension under dirs is
// modified.
//
// .css is watched at all for a reason: the build compiles Tailwind, and the entry
// stylesheet (its @theme block, its @font-face rules, any hand-written CSS) is an
// INPUT to that build. Watching only .go meant saving app.css did nothing — no
// rebuild, no reload — and the change simply never reached the browser. Editing some
// unrelated Go file would then sweep it up by accident, which is a maddening way to
// find out.
func fingerprintExt(dirs []string, ext string) int64 {
	var fp int64
	for _, d := range dirs {
		filepath.WalkDir(d, func(path string, e fs.DirEntry, err error) error {
			if err != nil || e.IsDir() || !strings.EqualFold(filepath.Ext(path), ext) {
				return nil
			}
			if info, err := e.Info(); err == nil {
				fp += info.ModTime().UnixNano() + info.Size()
			}
			return nil
		})
	}
	return fp
}

func jsonStr(s string) string { b, _ := json.Marshal(s); return string(b) }

// ensureShim copies Go's wasm_exec.js into dir if it isn't already there, so the
// server is self-sufficient without a separate build step first.
func ensureShim(dir string) error {
	dst := filepath.Join(dir, "wasm_exec.js")
	if _, err := os.Stat(dst); err == nil {
		return nil
	}
	root, err := exec.Command("go", "env", "GOROOT").Output()
	if err != nil {
		return err
	}
	goroot := strings.TrimSpace(string(root))
	for _, p := range []string{
		filepath.Join(goroot, "lib", "wasm", "wasm_exec.js"),  // Go >= 1.24
		filepath.Join(goroot, "misc", "wasm", "wasm_exec.js"), // Go <= 1.23
	} {
		if b, err := os.ReadFile(p); err == nil {
			if err := os.MkdirAll(dir, 0o755); err != nil {
				return err
			}
			return os.WriteFile(dst, b, 0o644)
		}
	}
	return os.ErrNotExist
}

// ---- injected livereload client -----------------------------------------

const clientJS = `// Injected by the dev server in watch mode.
(function () {
  // Hot-swap the freshly built wasm in place — no full page reload and no blank
  // flash. We fetch + compile the NEW module while the current page stays
  // visible, then tear down the old instance and start the new one in the SAME
  // task, so the browser never paints the intermediate empty #app. State is
  // preserved: the old instance snapshots its signals on dispose; the new one
  // restores them on boot.
  async function hotSwap() {
    if (!window.__gowasmPrepare) { location.reload(); return; } // fallback
    var start;
    try { start = await window.__gowasmPrepare(); } // network + compile; page still visible
    catch (err) { console.error(err); location.reload(); return; }
    hideOverlay();
    try { if (window.__gowasmDispose) window.__gowasmDispose(); } catch (err) { console.error(err); }
    start(); // renders synchronously — no await between dispose and first paint
  }

  // Swap every stylesheet for a freshly-fetched copy — no page reload, so scroll
  // position, form state and the running wasm instance all survive a style tweak.
  //
  // The NEW link is loaded and only then does the old one go, on its onload. Removing
  // it first would leave the page unstyled for however long the fetch takes, which
  // reads as a flash of naked HTML every time you save.
  function swapCSS() {
    hideOverlay();
    var links = document.querySelectorAll('link[rel="stylesheet"]');
    for (var i = 0; i < links.length; i++) {
      (function (old) {
        var url = old.href.split("?")[0] + "?v=" + Date.now();
        var next = old.cloneNode();
        next.href = url;
        next.onload = function () { if (old.parentNode) old.parentNode.removeChild(old); };
        next.onerror = function () { if (next.parentNode) next.parentNode.removeChild(next); };
        old.parentNode.insertBefore(next, old.nextSibling);
      })(links[i]);
    }
    console.log("[hot reload] stylesheet updated");
  }

  // Full-screen overlay showing the Go compiler output when a build fails. The
  // app underneath keeps running (and its state), so fixing the code and saving
  // clears the overlay and hot-swaps without losing anything.
  function ensureOverlay() {
    var el = document.getElementById("__gowasm_error");
    if (!el) {
      el = document.createElement("div");
      el.id = "__gowasm_error";
      el.style.cssText = "position:fixed;inset:0;z-index:2147483647;margin:0;padding:24px 28px;" +
        "background:rgba(24,24,27,0.97);color:#e4e4e7;overflow:auto;" +
        "font:13px/1.55 ui-monospace,SFMono-Regular,Menlo,Consolas,monospace;";
      document.body.appendChild(el);
    }
    return el;
  }
  function showOverlay(text) {
    var el = ensureOverlay();
    el.textContent = "";
    var head = document.createElement("div");
    head.style.cssText = "color:#f87171;font-weight:700;font-size:15px;margin-bottom:14px;";
    head.textContent = "⚠  gowasm — build failed";
    var pre = document.createElement("pre");
    pre.style.cssText = "margin:0;white-space:pre-wrap;word-break:break-word;";
    pre.textContent = text || "(no compiler output)";
    el.appendChild(head);
    el.appendChild(pre);
  }
  function hideOverlay() {
    var el = document.getElementById("__gowasm_error");
    if (el && el.parentNode) { el.parentNode.removeChild(el); }
  }
  window.__gowasmErrorOverlay = { show: showOverlay, hide: hideOverlay };

  function connect() {
    var proto = location.protocol === "https:" ? "wss://" : "ws://";
    var ws = new WebSocket(proto + location.host + "/livereload");
    ws.onmessage = function (e) {
      var msg = {};
      try { msg = JSON.parse(e.data); } catch (_) { return; }
      if (msg.type === "reload") { hotSwap(); }                 // build ok: swap in place
      else if (msg.type === "css") { swapCSS(); }               // styles only: swap the stylesheet
      else if (msg.type === "error") { showOverlay(msg.msg); }  // build failed: show compiler output
      else if (msg.type === "building") { console.log("[hot reload] rebuilding…"); }
    };
    ws.onclose = function () { setTimeout(connect, 1000); }; // reconnect after reload/restart
    ws.onerror = function () { try { ws.close(); } catch (_) {} };
  }
  connect();
})();
`

func serveClientJS(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/javascript; charset=utf-8")
	w.Header().Set("Cache-Control", "no-store")
	io.WriteString(w, clientJS)
}

// ---- minimal WebSocket hub (stdlib only) --------------------------------

type hub struct {
	mu        sync.Mutex
	clients   map[*wsConn]struct{}
	lastError string // most recent build-failure message (JSON), replayed to new clients
}

func newHub() *hub { return &hub{clients: map[*wsConn]struct{}{}} }

func (h *hub) add(c *wsConn)    { h.mu.Lock(); h.clients[c] = struct{}{}; h.mu.Unlock() }
func (h *hub) remove(c *wsConn) { h.mu.Lock(); delete(h.clients, c); h.mu.Unlock() }

// setError records the current build failure (so it survives to new clients) and
// pushes it to everyone connected. clearError is called on the next good build.
func (h *hub) setError(out string) {
	msg := `{"type":"error","msg":` + jsonStr(out) + `}`
	h.mu.Lock()
	h.lastError = msg
	h.mu.Unlock()
	h.broadcast(msg)
}

func (h *hub) clearError() { h.mu.Lock(); h.lastError = ""; h.mu.Unlock() }

func (h *hub) errorMessage() string { h.mu.Lock(); defer h.mu.Unlock(); return h.lastError }

func (h *hub) broadcast(msg string) {
	h.mu.Lock()
	defer h.mu.Unlock()
	for c := range h.clients {
		if err := c.sendText(msg); err != nil {
			c.conn.Close()
			delete(h.clients, c)
		}
	}
}

// serveWS upgrades the request to a WebSocket and keeps the connection until the
// client disconnects. We only ever push server->client, so incoming frames are
// drained (which also lets us detect disconnects).
func (h *hub) serveWS(w http.ResponseWriter, r *http.Request) {
	if !strings.Contains(strings.ToLower(r.Header.Get("Upgrade")), "websocket") {
		http.Error(w, "expected websocket upgrade", http.StatusBadRequest)
		return
	}
	hj, ok := w.(http.Hijacker)
	if !ok {
		http.Error(w, "hijacking unsupported", http.StatusInternalServerError)
		return
	}
	conn, brw, err := hj.Hijack()
	if err != nil {
		return
	}
	accept := acceptKey(r.Header.Get("Sec-WebSocket-Key"))
	io.WriteString(brw, "HTTP/1.1 101 Switching Protocols\r\n"+
		"Upgrade: websocket\r\nConnection: Upgrade\r\n"+
		"Sec-WebSocket-Accept: "+accept+"\r\n\r\n")
	if brw.Flush() != nil {
		conn.Close()
		return
	}

	c := &wsConn{conn: conn}
	h.add(c)
	if msg := h.errorMessage(); msg != "" {
		c.sendText(msg) // opened after a failed build => show the overlay right away
	}
	// Drain incoming bytes; return (and clean up) when the client goes away.
	go func() {
		io.Copy(io.Discard, brw)
		conn.Close()
		h.remove(c)
	}()
}

func acceptKey(key string) string {
	h := sha1.New()
	io.WriteString(h, key+"258EAFA5-E914-47DA-95CA-C5AB0DC85B11")
	return base64.StdEncoding.EncodeToString(h.Sum(nil))
}

type wsConn struct {
	mu   sync.Mutex
	conn net.Conn
}

// sendText writes a single unmasked text frame (server frames are never masked).
func (c *wsConn) sendText(msg string) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	payload := []byte(msg)
	n := len(payload)
	var header []byte
	switch {
	case n < 126:
		header = []byte{0x81, byte(n)}
	case n < 1<<16:
		header = []byte{0x81, 126, byte(n >> 8), byte(n)}
	default:
		header = []byte{0x81, 127,
			byte(n >> 56), byte(n >> 48), byte(n >> 40), byte(n >> 32),
			byte(n >> 24), byte(n >> 16), byte(n >> 8), byte(n)}
	}
	if _, err := c.conn.Write(header); err != nil {
		return err
	}
	_, err := c.conn.Write(payload)
	return err
}