package mill

import (
	"net/http"
	"strings"

	"github.com/gorilla/websocket"

	millproto "tangled.org/core/spindle/mill/proto"
	millv1 "tangled.org/core/spindle/mill/proto/gen"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

// HandleExecutorConn terminates one executor websocket. Auth is a pre-shared
// secret in the Authorization header, checked before the upgrade so a bad token
// never opens a socket.
func (m *Mill) HandleExecutorConn(w http.ResponseWriter, r *http.Request) {
	name, ok := m.authenticate(r)
	if !ok {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		m.l.Error("fleet ws upgrade failed", "err", err)
		return
	}
	defer conn.Close()

	stream := millproto.NewWSStream(conn)
	enc := millproto.NewEncoder(stream)
	dec := millproto.NewDecoder(stream)

	// first frame must be Hello.
	hello, err := dec.Decode()
	if err != nil {
		m.l.Error("fleet read hello failed", "err", err)
		return
	}
	h := hello.GetHello()
	if h == nil {
		m.l.Error("fleet first frame was not hello")
		return
	}
	if h.GetProtocolVersion() != millproto.ProtocolVersion {
		m.l.Error("fleet protocol version mismatch", "got", h.GetProtocolVersion(), "want", millproto.ProtocolVersion)
		return
	}

	// identity comes from the authenticated token, never the client-supplied
	// Hello node id, so a valid token can't impersonate another executor.
	sess := newSession(name, enc, m.l)
	sess.closeTransport = conn.Close
	sess.labels = h.GetLabels()
	resume, ok := m.attachSession(sess)
	if !ok {
		m.l.Warn("rejecting duplicate live executor session", "node", name)
		return
	}
	m.l.Info("executor connected", "node", sess.nodeID, "arch", h.GetArch(), "labels", h.GetLabels(), "resume", resume)

	if err := sess.send(&millproto.Message{Resume: &millv1.Resume{AckOffset: resume}}); err != nil {
		m.l.Error("fleet send resume failed", "err", err)
		m.detachSession(sess)
		return
	}
	m.sessionReady(sess)

	if err := sess.readLoop(m, dec); err != nil {
		m.l.Debug("session read ended", "node", sess.nodeID, "err", err)
	}
	m.detachSession(sess)
}

// authenticate resolves the executor identity from the pre-shared token in the
// Authorization header before the upgrade, so a bad token never opens a socket.
// The token is looked up by hash; an unknown or missing token is rejected
// (fail closed).
func (m *Mill) authenticate(r *http.Request) (string, bool) {
	const prefix = "Bearer "
	h := r.Header.Get("Authorization")
	if !strings.HasPrefix(h, prefix) {
		return "", false
	}
	token := strings.TrimPrefix(h, prefix)
	if token == "" || m.db == nil {
		return "", false
	}
	name, ok, err := m.db.ResolveExecutorToken(HashToken(token))
	if err != nil {
		m.l.Error("executor token lookup failed", "err", err)
		return "", false
	}
	return name, ok
}
