package mill

import (
	"context"
	"errors"
	"log/slog"
	"sync"
	"time"

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

var errSessionClosed = errors.New("mill: executor session closed")

// millSession is one live websocket to an executor. Unlike the microVM
// AgentSession it is multiplexed: many leases plus async status/log relay share
// it, so a single reader goroutine demuxes by message type and correlates
// request/response by lease id. We never hold a lock across a decode.
type millSession struct {
	nodeID         string
	labels         []string
	enc            messageEncoder
	l              *slog.Logger
	closeTransport func() error

	// snapshot, disconnected and graceTimer are guarded by Mill.mu, not the
	// session mutex below (the fleet ranks across sessions under its own lock).
	snapshot     *millv1.NodeSnapshot
	disconnected bool
	graceTimer   *time.Timer
	lastSeen     time.Time

	mu      sync.Mutex
	pending map[string]chan *millproto.Message // lease_id -> response waiter

	closeOnce sync.Once
	closed    chan struct{}
}

type messageEncoder interface {
	Encode(*millproto.Message) error
}

func newSession(nodeID string, enc messageEncoder, l *slog.Logger) *millSession {
	return &millSession{
		nodeID:   nodeID,
		enc:      enc,
		l:        l,
		pending:  make(map[string]chan *millproto.Message),
		closed:   make(chan struct{}),
		lastSeen: time.Now(),
	}
}

func (s *millSession) send(msg *millproto.Message) error {
	return s.enc.Encode(msg)
}

func (s *millSession) close() {
	s.closeOnce.Do(func() {
		close(s.closed)
		if s.closeTransport != nil {
			_ = s.closeTransport()
		}
	})
}

// await registers a one-shot waiter for the next response addressed to leaseID,
// returning the channel and a cancel that unregisters it.
func (s *millSession) await(leaseID string) (<-chan *millproto.Message, func()) {
	ch := make(chan *millproto.Message, 1)
	s.mu.Lock()
	s.pending[leaseID] = ch
	s.mu.Unlock()
	return ch, func() {
		s.mu.Lock()
		if s.pending[leaseID] == ch {
			delete(s.pending, leaseID)
		}
		s.mu.Unlock()
	}
}

// deliver routes a correlated response to a waiter, if any is registered.
func (s *millSession) deliver(leaseID string, msg *millproto.Message) {
	s.mu.Lock()
	ch := s.pending[leaseID]
	delete(s.pending, leaseID)
	s.mu.Unlock()
	if ch != nil {
		select {
		case ch <- msg:
		default:
		}
	}
}

// request sends msg and waits for the correlated response addressed to leaseID,
// honouring ctx and session close.
func (s *millSession) request(ctx context.Context, leaseID string, msg *millproto.Message) (*millproto.Message, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	ch, cancel := s.await(leaseID)
	defer cancel()

	if err := s.send(msg); err != nil {
		return nil, err
	}

	select {
	case resp := <-ch:
		return resp, nil
	case <-ctx.Done():
		return nil, ctx.Err()
	case <-s.closed:
		return nil, errSessionClosed
	}
}

// readLoop demuxes incoming frames until the decoder errors (connection gone).
func (s *millSession) readLoop(m *Mill, dec *millproto.Decoder) error {
	for {
		msg, err := dec.Decode()
		if err != nil {
			return err
		}
		if err := s.dispatch(m, msg); err != nil {
			return err
		}
	}
}

func (s *millSession) dispatch(m *Mill, msg *millproto.Message) error {
	if !m.touchSession(s) {
		return errSessionClosed
	}
	switch {
	case msg.GetNodeSnapshot() != nil:
		return m.onSnapshot(s, msg.GetNodeSnapshot())
	case msg.GetReserveResult() != nil:
		s.deliver(msg.GetReserveResult().GetLeaseId(), msg)
	case msg.GetCommitted() != nil:
		s.deliver(msg.GetCommitted().GetLeaseId(), msg)
	case msg.GetStatusEvent() != nil:
		return m.onStatusRelay(s, msg.GetStatusEvent())
	case msg.GetLogLine() != nil:
		return m.onLogRelay(s, msg.GetLogLine())
	case msg.GetAttemptResult() != nil:
		return m.onAttemptResult(s, msg.GetAttemptResult())
	default:
		s.l.Warn("session received unexpected message", "node", s.nodeID)
	}
	return nil
}
