package main

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"html"
	"io"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"

	"wxch.org/wxch/internal/modules"
)

// A tunnel lets the server show an engine's own web page (ComfyUI, for the administrator) without opening any port on
// this machine: the server sends the browser's request over the agent connection, the agent makes it to the engine on
// 127.0.0.1 and sends the answer back. HTTP requests and websockets are both carried.
//
//	server -> agent   tun_req (Method=engine, D=path, Params={method,headers}), tun_body (D=base64, Status=1 at the end),
//	                  tun_ws_open (Method=engine, D=path), tun_ws_msg (D=base64, Status=opcode), tun_ws_close, tun_cancel
//	agent  -> server  tun_res (Status=http code, Params=headers), tun_data (D=base64), tun_end (Err),
//	                  tun_ws_msg, tun_ws_close

const tunChunk = 192 << 10

type agentTun struct {
	in     chan Msg
	cancel context.CancelFunc
}

var tuns = struct {
	sync.Mutex
	m map[string]*agentTun
}{m: map[string]*agentTun{}}

// tunnelsOpen says whether anyone is using an engine's page; the idle stop must leave it alone.
func tunnelsOpen() bool {
	tuns.Lock()
	defer tuns.Unlock()
	return len(tuns.m) > 0
}

func touchEngine(engine string) {
	jobs.Lock()
	jobs.lastUse[engine] = time.Now()
	jobs.Unlock()
}

// handleTunMsg is called from the connection's read loop, in order; the work itself runs in goroutines.
func handleTunMsg(m Msg, send func(Msg)) {
	switch m.T {
	case "tun_req", "tun_ws_open":
		ctx, cancel := context.WithCancel(context.Background())
		t := &agentTun{in: make(chan Msg, 256), cancel: cancel}
		tuns.Lock()
		tuns.m[m.ID] = t
		tuns.Unlock()
		if m.T == "tun_req" {
			go runTunHTTP(ctx, m, t, send)
		} else {
			go runTunWS(ctx, m, t, send)
		}
	default:
		tuns.Lock()
		t := tuns.m[m.ID]
		tuns.Unlock()
		if t == nil {
			return
		}
		if m.T == "tun_cancel" {
			t.cancel()
			return
		}
		select {
		case t.in <- m:
		case <-time.After(10 * time.Second):
			t.cancel() // the engine side is not keeping up
		}
	}
}

func closeTun(id string) {
	tuns.Lock()
	if t := tuns.m[id]; t != nil {
		t.cancel()
		delete(tuns.m, id)
	}
	tuns.Unlock()
}

// tunTarget resolves an engine to its local address and tells whether it answers. An engine that is not running is
// started in the background and the caller shows a "starting" page.
func tunTarget(engine string) (hostport, base string, ready bool, err error) {
	cfg := loadNodeConfig()
	cat, err := loadCatalog()
	if err != nil {
		return "", "", false, err
	}
	e, ok := cat.Engine(engine)
	if !ok {
		return "", "", false, fmt.Errorf("unknown engine")
	}
	if len(capabilitiesUsing(cat, readModuleState(cfg), engine)) == 0 {
		return "", "", false, fmt.Errorf("no module that uses this engine is installed on this machine")
	}
	base = engineBase(cat, e)
	u, _ := url.Parse(base)
	hostport = u.Host
	hc := &http.Client{Timeout: 2 * time.Second}
	if resp, err := hc.Get(base + e.Health); err == nil {
		resp.Body.Close()
		if resp.StatusCode == 200 {
			return hostport, base, true, nil
		}
	}
	touchEngine(e.ID) // someone is waiting for it: it is not idle
	maybeStartEngine(cfg, cat, e)
	return hostport, base, false, nil
}

var lastEngineStart sync.Map // engine -> time of the last start we asked for

// maybeStartEngine starts the engine when it is not running, but not again within a minute of the last attempt: an engine
// that dies at start-up (a package that cannot be installed) must not be restarted on every reload of the page.
func maybeStartEngine(cfg nodeConfig, cat *modules.Catalog, e modules.Engine) {
	if rn, err := pickModRunner(cfg); err == nil {
		if st, err := rn.Status(); err == nil {
			for _, s := range st {
				if s.Engine == e.ID && s.State == "running" {
					return // it is starting up: the page shows its log
				}
			}
		}
	}
	if t, ok := lastEngineStart.Load(e.ID); ok && time.Since(t.(time.Time)) < time.Minute {
		return
	}
	lastEngineStart.Store(e.ID, time.Now())
	go startEngineQuietly(cfg, cat, e)
}

// engineState says in words what the engine container is doing, and its last log lines.
func engineState(engine string) (state, logs string) {
	cfg := loadNodeConfig()
	rn, err := pickModRunner(cfg)
	if err != nil {
		return "", err.Error()
	}
	if st, err := rn.Status(); err == nil {
		for _, s := range st {
			if s.Engine == engine {
				state = s.State
				if state == "" && s.Pulling {
					state = "pulling the image"
				}
			}
		}
	}
	return state, rn.Logs(engine, 25)
}

var startingEngine sync.Map

func startEngineQuietly(cfg nodeConfig, cat *modules.Catalog, e modules.Engine) {
	if _, busy := startingEngine.LoadOrStore(e.ID, true); busy {
		return
	}
	defer startingEngine.Delete(e.ID)
	ctx, cancel := context.WithTimeout(context.Background(), engineWait)
	defer cancel()
	if _, err := ensureEngine(ctx, cfg, cat, e, func(string, float64) {}); err != nil {
		logf("engine %s: %v", e.ID, err)
		lastEngineErr.Store(e.ID, err.Error()) // the starting page shows why
	} else {
		lastEngineErr.Delete(e.ID)
		touchEngine(e.ID) // it is ready now: the idle clock starts here
	}
}

var lastEngineErr sync.Map // engine -> why the last start failed

const startingPage = `<!doctype html><meta charset="utf-8"><meta http-equiv="refresh" content="5"><title>starting</title>
<body style="font:16px system-ui;margin:3em"><h2>กำลังเริ่ม engine… / Starting the engine…</h2><p>หน้านี้จะโหลดซ้ำเองทุก 5 วินาที (ครั้งแรกอาจใช้เวลาหลายนาที) / This page reloads every 5 seconds (the first start can take minutes).</p>`

// startingPageFor is the page shown while the engine is not ready: it reloads itself and shows what the engine says.
func startingPageFor(engine string) string {
	state, logs := engineState(engine)
	if state == "" {
		state = "starting"
	}
	page := startingPage
	if why, ok := lastEngineErr.Load(engine); ok {
		page += "<p style=\"color:#c33\"><b>ไม่สามารถเริ่ม engine ได้ / The engine could not be started:</b> " + html.EscapeString(why.(string)) + "</p>"
	}
	return page + "<p><b>" + html.EscapeString(state) + "</b></p><pre style=\"white-space:pre-wrap;font-size:12px;opacity:.8;max-height:60vh;overflow:auto\">" + html.EscapeString(logs) + "</pre>"
}

func send503(send func(Msg), id, body string) {
	h, _ := json.Marshal(map[string]string{"Content-Type": "text/html; charset=utf-8", "Cache-Control": "no-store"})
	send(Msg{T: "tun_res", ID: id, Status: 503, Params: h})
	sendChunks(send, id, []byte(body))
	send(Msg{T: "tun_end", ID: id})
}

func sendChunks(send func(Msg), id string, b []byte) {
	for len(b) > 0 {
		n := len(b)
		if n > tunChunk {
			n = tunChunk
		}
		send(Msg{T: "tun_data", ID: id, D: base64.StdEncoding.EncodeToString(b[:n])})
		b = b[n:]
	}
}

var hopHeaders = map[string]bool{"connection": true, "keep-alive": true, "proxy-authenticate": true, "proxy-authorization": true, "te": true,
	"trailer": true, "transfer-encoding": true, "upgrade": true, "content-length": true}

func runTunHTTP(ctx context.Context, m Msg, t *agentTun, send func(Msg)) {
	defer closeTun(m.ID)
	var p struct {
		Method  string            `json:"method"`
		Headers map[string]string `json:"headers"`
	}
	json.Unmarshal(m.Params, &p)
	if m.Method == fetchEngine && strings.HasPrefix(m.D, "/_wxch/") { // for the agent itself, not for the engine
		serveLocalTun(ctx, m, t, p.Method, send)
		return
	}
	hostport, base, ready, err := tunTarget(m.Method)
	if err != nil {
		send(Msg{T: "tun_res", ID: m.ID, Status: 502})
		send(Msg{T: "tun_end", ID: m.ID, Err: err.Error()})
		return
	}
	if !ready {
		send503(send, m.ID, startingPageFor(m.Method))
		return
	}
	touchEngine(m.Method)
	defer touchEngine(m.Method)
	pr, pw := io.Pipe()
	var body io.Reader
	hasBody := p.Method != "GET" && p.Method != "HEAD"
	if hasBody {
		body = pr
		go func() { // the request body arrives as tun_body messages
			defer pw.Close()
			for {
				select {
				case <-ctx.Done():
					return
				case b := <-t.in:
					if b.T != "tun_body" {
						continue
					}
					if raw, err := base64.StdEncoding.DecodeString(b.D); err == nil && len(raw) > 0 {
						if _, err := pw.Write(raw); err != nil {
							return
						}
					}
					if b.Status == 1 {
						return
					}
				}
			}
		}()
	}
	req, err := http.NewRequestWithContext(ctx, p.Method, base+m.D, body)
	if err != nil {
		send(Msg{T: "tun_res", ID: m.ID, Status: 400})
		send(Msg{T: "tun_end", ID: m.ID, Err: "bad request"})
		return
	}
	for k, v := range p.Headers {
		if hopHeaders[strings.ToLower(k)] || strings.EqualFold(k, "host") || strings.EqualFold(k, "cookie") || strings.EqualFold(k, "authorization") {
			continue
		}
		req.Header.Set(k, v)
	}
	// the engine only accepts a page whose Origin matches its own address (its protection against other sites)
	req.Host = hostport
	if req.Header.Get("Origin") != "" {
		req.Header.Set("Origin", "http://"+hostport)
	}
	req.Header.Set("Referer", "http://"+hostport+"/")
	if cl := p.Headers["Content-Length"]; cl != "" && hasBody {
		var n int64
		fmt.Sscan(cl, &n)
		req.ContentLength = n
	}
	resp, err := (&http.Client{Transport: &http.Transport{DisableCompression: true}, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}).Do(req)
	if err != nil {
		send(Msg{T: "tun_res", ID: m.ID, Status: 502})
		send(Msg{T: "tun_end", ID: m.ID, Err: err.Error()})
		return
	}
	defer resp.Body.Close()
	hdr := map[string]string{}
	for k, v := range resp.Header {
		if !hopHeaders[strings.ToLower(k)] && len(v) > 0 {
			hdr[k] = v[0]
		}
	}
	h, _ := json.Marshal(hdr)
	send(Msg{T: "tun_res", ID: m.ID, Status: resp.StatusCode, Params: h})
	buf := make([]byte, tunChunk)
	for {
		n, rerr := resp.Body.Read(buf)
		if n > 0 {
			send(Msg{T: "tun_data", ID: m.ID, D: base64.StdEncoding.EncodeToString(buf[:n])})
		}
		if rerr != nil {
			send(Msg{T: "tun_end", ID: m.ID})
			return
		}
		select {
		case <-ctx.Done():
			return
		default:
		}
	}
}

func runTunWS(ctx context.Context, m Msg, t *agentTun, send func(Msg)) {
	defer closeTun(m.ID)
	hostport, base, ready, err := tunTarget(m.Method)
	if err != nil || !ready {
		send(Msg{T: "tun_ws_close", ID: m.ID})
		return
	}
	hdr := http.Header{"Origin": {"http://" + hostport}}
	ws, _, err := websocket.DefaultDialer.DialContext(ctx, "ws"+strings.TrimPrefix(base, "http")+m.D, hdr)
	if err != nil {
		send(Msg{T: "tun_ws_close", ID: m.ID})
		return
	}
	defer ws.Close()
	ws.SetReadLimit(32 << 20)
	touchEngine(m.Method)
	go func() { // keeps the engine from being stopped as idle while the page is open
		tk := time.NewTicker(20 * time.Second)
		defer tk.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-tk.C:
				touchEngine(m.Method)
			}
		}
	}()
	go func() {
		defer closeTun(m.ID)
		for {
			mt, data, err := ws.ReadMessage()
			if err != nil {
				send(Msg{T: "tun_ws_close", ID: m.ID})
				return
			}
			send(Msg{T: "tun_ws_msg", ID: m.ID, Status: mt, D: base64.StdEncoding.EncodeToString(data)})
		}
	}()
	for {
		select {
		case <-ctx.Done():
			return
		case b := <-t.in:
			switch b.T {
			case "tun_ws_msg":
				raw, err := base64.StdEncoding.DecodeString(b.D)
				if err != nil {
					continue
				}
				mt := b.Status
				if mt != websocket.BinaryMessage {
					mt = websocket.TextMessage
				}
				if ws.WriteMessage(mt, raw) != nil {
					return
				}
			case "tun_ws_close":
				return
			}
		}
	}
}