package main

import (
	"encoding/base64"
	"encoding/json"
	"html"
	"io"
	"log"
	"net/http"
	"regexp"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	"golang.org/x/crypto/bcrypt"
)

// comfy.wxch.org: ComfyUI's own web page, for the administrator. The browser talks to this server; the server carries the
// requests to the machine over its agent connection (see cmd/wxch-agent/tunnel.go), so no port is opened anywhere and the
// engine starts when the page is opened. ComfyUI can run any workflow, so the page needs the web password (and the 2FA
// code when one is configured), and it is not linked from anything public.

const (
	comfyEngine   = "comfyui"
	tunBodyChunk  = 192 << 10
	maxTunBody    = 1 << 30
	tunWaitResult = 2 * time.Minute
)

type srvTun struct {
	ch    chan Msg
	agent string
	done  chan struct{}
	once  sync.Once
}

type tunRegistry struct {
	mu sync.Mutex
	m  map[string]*srvTun
}

func (r *tunRegistry) open(id, agent string) *srvTun {
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.m == nil {
		r.m = map[string]*srvTun{}
	}
	t := &srvTun{ch: make(chan Msg, 256), agent: agent, done: make(chan struct{})}
	r.m[id] = t
	return t
}

func (r *tunRegistry) close(id string) {
	r.mu.Lock()
	t := r.m[id]
	delete(r.m, id)
	r.mu.Unlock()
	if t != nil {
		t.once.Do(func() { close(t.done) })
	}
}

// deliver passes a message from the agent to the request waiting for it.
func (r *tunRegistry) deliver(agent string, m Msg) {
	r.mu.Lock()
	t := r.m[m.ID]
	r.mu.Unlock()
	if t == nil || t.agent != agent {
		return
	}
	select {
	case t.ch <- m:
	case <-t.done:
	case <-time.After(10 * time.Second):
	}
}

// agentGone ends every tunnel of a machine that disconnected.
func (r *tunRegistry) agentGone(agent string) {
	r.mu.Lock()
	var ids []string
	for id, t := range r.m {
		if t.agent == agent {
			ids = append(ids, id)
		}
	}
	r.mu.Unlock()
	for _, id := range ids {
		r.mu.Lock()
		t := r.m[id]
		r.mu.Unlock()
		if t != nil {
			select {
			case t.ch <- Msg{T: "tun_end", ID: id, Err: "machine disconnected"}:
			default:
			}
		}
	}
}

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

// comfyMachine picks the machine whose ComfyUI is shown: the one asked for with ?machine=, else the first online machine
// that has a module running on ComfyUI.
func (s *Server) comfyMachine(r *http.Request) (string, *conn) {
	h := s.hub
	h.mu.Lock()
	defer h.mu.Unlock()
	var names []string
	for n := range h.agents {
		b := h.stats[n]
		if b == nil || time.Since(b.at) > 60*time.Second {
			continue
		}
		for id := range b.modules {
			if cp, ok := s.jobs.cat.Capability(id); ok && cp.Engine == comfyEngine {
				names = append(names, n)
				break
			}
		}
	}
	if len(names) == 0 {
		return "", nil
	}
	sort.Strings(names)
	want := r.URL.Query().Get("machine")
	if want == "" {
		if c, err := r.Cookie("wxch_comfy_machine"); err == nil {
			want = c.Value
		}
	}
	for _, n := range names {
		if n == want {
			return n, h.agents[n]
		}
	}
	return names[0], h.agents[names[0]]
}

const comfyLoginPage = `<!doctype html><html lang="th"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1"><meta name="robots" content="noindex"><title>ComfyUI</title><link rel="icon" type="image/svg+xml" href="data:image/svg+xml,%3Csvg%20xmlns%3D%22http%3A%2F%2Fwww.w3.org%2F2000%2Fsvg%22%20viewBox%3D%220%200%2032%2032%22%3E%3Cpolygon%20points%3D%2216%2C3%2028%2C10%2016%2C17%204%2C10%22%20fill%3D%22%238fe8b8%22%2F%3E%3Cpolygon%20points%3D%224%2C10%2016%2C17%2016%2C30%204%2C23%22%20fill%3D%22%233fae7a%22%2F%3E%3Cpolygon%20points%3D%2216%2C17%2028%2C10%2028%2C23%2016%2C30%22%20fill%3D%22%232a8a5c%22%2F%3E%3C%2Fsvg%3E">
<style>:root{color-scheme:light dark}body{font-family:system-ui,"Leelawadee UI",sans-serif;margin:0;min-height:100vh;display:grid;place-items:center}form{width:min(340px,90vw);display:grid;gap:10px}
input,button{font:inherit;padding:9px 11px;border-radius:8px;border:1px solid #8886}button{cursor:pointer;font-weight:600}#e{color:#c33;min-height:1.3em}</style></head>
<body><form id="f"><h2 style="margin:0"><svg style="vertical-align:-5px;margin-right:6px" viewBox="0 0 32 32" width="24" height="24" aria-hidden="true"><polygon points="16,3 28,10 16,17 4,10" fill="#8fe8b8"/><polygon points="4,10 16,17 16,30 4,23" fill="#3fae7a"/><polygon points="16,17 28,10 28,23 16,30" fill="#2a8a5c"/></svg>ComfyUI (admin)</h2><input id="pw" type="password" placeholder="รหัสผ่าน / Password" autocomplete="current-password" required>
<input id="cd" inputmode="numeric" placeholder="รหัส 2FA (ถ้ามี) / 2FA code (if set)" autocomplete="one-time-code"><button>เข้า / Enter</button><div id="e"></div></form>
<script>f.onsubmit=async e=>{e.preventDefault();e.target.querySelector('button').disabled=true;const r=await fetch('/comfy-login',{method:'POST',body:JSON.stringify({Password:pw.value,Code:cd.value})});
if(r.ok)location.reload();else{document.getElementById('e').textContent=r.status==429?'ลองผิดบ่อยเกินไป รอสักครู่ / Too many attempts':'ไม่ถูกต้อง / Invalid';e.target.querySelector('button').disabled=false}}</script></body></html>`

func (s *Server) comfyMux() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/comfy-login" && r.Method == http.MethodPost {
			s.comfyLogin(w, r)
			return
		}
		if !s.comfyAuthed(r) {
			if r.Method == http.MethodGet && strings.Contains(r.Header.Get("Accept"), "text/html") {
				w.Header().Set("Content-Type", "text/html; charset=utf-8")
				w.Header().Set("Cache-Control", "no-store")
				w.Write([]byte(comfyLoginPage))
				return
			}
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		name, c := s.comfyMachine(r)
		if c == nil {
			comfyMessage(w, r, http.StatusServiceUnavailable, "ยังไม่มีเครื่องที่ติดตั้งโมดูลภาพ/วิดีโอ", "No machine has an image or video module installed",
				"ติดตั้งได้ที่หน้า \"โมดูล AI\" ของ control / Install one on the AI modules page of the control site")
			return
		}
		if q := r.URL.Query().Get("machine"); q == name {
			http.SetCookie(w, &http.Cookie{Name: "wxch_comfy_machine", Value: name, Path: "/", HttpOnly: true, Secure: true, SameSite: http.SameSiteStrictMode})
		}
		if websocket.IsWebSocketUpgrade(r) {
			s.comfyWS(w, r, name, c)
			return
		}
		s.comfyHTTP(w, r, name, c)
	})
}

func (s *Server) comfyLogin(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, Code string }
	if json.NewDecoder(http.MaxBytesReader(w, r.Body, 4096)).Decode(&req) != nil {
		http.Error(w, "bad request", 400)
		return
	}
	pwOK := bcrypt.CompareHashAndPassword([]byte(s.pwHash()), []byte(req.Password)) == nil
	codeOK := s.cfg.TOTPSecret == "" || totpOK(s.cfg.TOTPSecret, req.Code)
	if !pwOK || !codeOK {
		s.fail(ip)
		log.Printf("comfy login failed from %s", ip)
		time.Sleep(time.Second)
		http.Error(w, "invalid", http.StatusUnauthorized)
		return
	}
	tok := randHex(32)
	s.comfySess.add(tok, 12*time.Hour)
	http.SetCookie(w, &http.Cookie{Name: "wxch_comfy", Value: tok, Path: "/", HttpOnly: true, Secure: true, SameSite: http.SameSiteStrictMode, MaxAge: 12 * 3600})
	log.Printf("comfy login ok from %s", ip)
	w.Write([]byte("ok"))
}

var tunSkipHeaders = map[string]bool{"connection": true, "keep-alive": true, "te": true, "trailer": true, "transfer-encoding": true, "upgrade": true,
	"cookie": true, "authorization": true, "accept-encoding": true, "host": true, "x-forwarded-for": true, "x-forwarded-proto": true,
	"x-forwarded-host": true, "x-real-ip": true, "cdn-loop": true, "forwarded": true}

func (s *Server) comfyHTTP(w http.ResponseWriter, r *http.Request, agent string, c *conn) {
	id := randHex(8)
	t := s.tuns.open(id, agent)
	defer s.tuns.close(id)
	hdr := map[string]string{}
	for k, v := range r.Header {
		if !tunSkipHeaders[strings.ToLower(k)] && len(v) > 0 {
			hdr[k] = v[0]
		}
	}
	if r.ContentLength > 0 {
		hdr["Content-Length"] = itoa(r.ContentLength)
	}
	params, _ := json.Marshal(map[string]any{"method": r.Method, "headers": hdr})
	c.send(Msg{T: "tun_req", ID: id, Method: comfyEngine, D: r.URL.RequestURI(), Params: params})
	if r.Method != http.MethodGet && r.Method != http.MethodHead && r.Body != nil {
		buf := make([]byte, tunBodyChunk)
		body := http.MaxBytesReader(w, r.Body, maxTunBody)
		for {
			n, err := body.Read(buf)
			if n > 0 {
				c.send(Msg{T: "tun_body", ID: id, D: base64.StdEncoding.EncodeToString(buf[:n])})
			}
			if err != nil {
				if err != io.EOF {
					c.send(Msg{T: "tun_cancel", ID: id})
					http.Error(w, "upload failed", http.StatusBadRequest)
					return
				}
				break
			}
		}
		c.send(Msg{T: "tun_body", ID: id, Status: 1})
	}
	fl, _ := w.(http.Flusher)
	started := false
	var page []byte // an HTML answer is held back so the logo can be added to it
	isPage, pageStatus := false, 0
	timeout := time.NewTimer(tunWaitResult)
	defer timeout.Stop()
	for {
		select {
		case m := <-t.ch:
			timeout.Reset(tunWaitResult)
			switch m.T {
			case "tun_res":
				var h map[string]string
				json.Unmarshal(m.Params, &h)
				for k, v := range h {
					w.Header().Set(k, v)
				}
				if strings.HasPrefix(strings.ToLower(h["Content-Type"]), "text/html") {
					isPage, pageStatus = true, m.Status // sent when it is complete
					w.Header().Del("Content-Length")
					w.Header().Del("Etag")
					w.Header().Del("Last-Modified")
					started = true
					continue
				}
				w.WriteHeader(m.Status)
				started = true
			case "tun_data":
				if raw, err := base64.StdEncoding.DecodeString(m.D); err == nil {
					if isPage {
						if len(page)+len(raw) <= 4<<20 {
							page = append(page, raw...)
						}
						continue
					}
					if _, err := w.Write(raw); err != nil {
						c.send(Msg{T: "tun_cancel", ID: id})
						return
					}
					if fl != nil {
						fl.Flush()
					}
				}
			case "tun_end":
				if isPage {
					out := brandHTML(page, r.Host)
					w.Header().Set("Content-Length", itoa(int64(len(out))))
					w.Header().Set("Cache-Control", "no-store")
					w.WriteHeader(pageStatus)
					w.Write(out)
					return
				}
				if !started {
					comfyMessage(w, r, http.StatusBadGateway, "เครื่องตอบไม่ได้", "The machine could not answer", m.Err)
				}
				return
			}
		case <-r.Context().Done():
			c.send(Msg{T: "tun_cancel", ID: id})
			return
		case <-timeout.C:
			c.send(Msg{T: "tun_cancel", ID: id})
			if !started {
				comfyMessage(w, r, http.StatusGatewayTimeout, "เครื่องไม่ตอบภายในเวลา", "The machine did not answer in time", "agent 0.12.2+ / agent 0.12.2 or newer")
			}
			return
		}
	}
}

// comfyMessage is a page of ours (a message to the administrator), with the same logo as the others.
func comfyMessage(w http.ResponseWriter, r *http.Request, code int, th, en, detail string) {
	body := brandHTML([]byte(`<!doctype html><html lang="th"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1"><meta name="robots" content="noindex"><title>ComfyUI</title>
<style>:root{color-scheme:light dark}body{font-family:system-ui,"Leelawadee UI",sans-serif;margin:0;min-height:100vh;display:grid;place-items:center;padding:24px;line-height:1.6}main{max-width:36em}h2{margin:0 0 .3em}p{opacity:.75;margin:.2em 0}</style></head>
<body><main><h2>`+html.EscapeString(th)+" / "+html.EscapeString(en)+"</h2><p>"+html.EscapeString(detail)+`</p></main></body></html>`), r.Host)
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	w.Header().Set("Cache-Control", "no-store")
	w.WriteHeader(code)
	w.Write(body)
}

// The page of ComfyUI is its own; the logo says whose it is: our icon replaces the page's, and a small mark in the corner
// links back to the studio.
var iconLinkRe = regexp.MustCompile(`(?is)<link[^>]+rel=["']?(?:shortcut )?icon["']?[^>]*>`)

const brandIcon = `<link rel="icon" type="image/svg+xml" href="data:image/svg+xml,%3Csvg%20xmlns%3D%22http%3A%2F%2Fwww.w3.org%2F2000%2Fsvg%22%20viewBox%3D%220%200%2032%2032%22%3E%3Cpolygon%20points%3D%2216%2C3%2028%2C10%2016%2C17%204%2C10%22%20fill%3D%22%238fe8b8%22%2F%3E%3Cpolygon%20points%3D%224%2C10%2016%2C17%2016%2C30%204%2C23%22%20fill%3D%22%233fae7a%22%2F%3E%3Cpolygon%20points%3D%2216%2C17%2028%2C10%2028%2C23%2016%2C30%22%20fill%3D%22%232a8a5c%22%2F%3E%3C%2Fsvg%3E">`

func brandHTML(page []byte, host string) []byte {
	s := iconLinkRe.ReplaceAllString(string(page), "")
	studio := "https://" + strings.Replace(host, "comfy.", "studio.", 1) + "/"
	badge := `<a href="` + studio + `" title="wxch studio" style="position:fixed;left:8px;bottom:8px;z-index:2147483000;opacity:.85;line-height:0"><svg viewBox="0 0 32 32" width="26" height="26" aria-hidden="true"><polygon points="16,3 28,10 16,17 4,10" fill="#8fe8b8"/><polygon points="4,10 16,17 16,30 4,23" fill="#3fae7a"/><polygon points="16,17 28,10 28,23 16,30" fill="#2a8a5c"/></svg></a>`
	if i := strings.Index(strings.ToLower(s), "</head>"); i >= 0 {
		s = s[:i] + brandIcon + s[i:]
	} else {
		s = brandIcon + s
	}
	badge += comfyFetchScript
	if i := strings.LastIndex(strings.ToLower(s), "</body>"); i >= 0 {
		s = s[:i] + badge + s[i:]
	} else {
		s += badge
	}
	return []byte(s)
}

func itoa(n int64) string { b, _ := json.Marshal(n); return string(b) }

func (s *Server) comfyWS(w http.ResponseWriter, r *http.Request, agent string, c *conn) {
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return
	}
	defer ws.Close()
	ws.SetReadLimit(16 << 20)
	id := randHex(8)
	t := s.tuns.open(id, agent)
	defer s.tuns.close(id)
	c.send(Msg{T: "tun_ws_open", ID: id, Method: comfyEngine, D: r.URL.RequestURI()})
	go func() { // browser -> machine
		defer func() { c.send(Msg{T: "tun_ws_close", ID: id}); s.tuns.close(id) }()
		for {
			mt, data, err := ws.ReadMessage()
			if err != nil {
				return
			}
			c.send(Msg{T: "tun_ws_msg", ID: id, Status: mt, D: base64.StdEncoding.EncodeToString(data)})
		}
	}()
	for { // machine -> browser
		select {
		case m := <-t.ch:
			switch m.T {
			case "tun_ws_msg":
				raw, err := base64.StdEncoding.DecodeString(m.D)
				if err != nil {
					continue
				}
				mt := m.Status
				if mt != websocket.BinaryMessage {
					mt = websocket.TextMessage
				}
				ws.SetWriteDeadline(time.Now().Add(15 * time.Second))
				if ws.WriteMessage(mt, raw) != nil {
					return
				}
			case "tun_ws_close", "tun_end":
				return
			}
		case <-t.done:
			return
		}
	}
}