package mill

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"slices"
	"strings"
	"sync"
	"time"

	"tangled.org/core/notifier"
	"tangled.org/core/spindle/db"
	"tangled.org/core/spindle/engine"
	"tangled.org/core/spindle/models"
	"tangled.org/core/spindle/secrets"
	"tangled.org/core/tid"

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

const (
	defaultReconnectGrace = 45 * time.Second
	defaultJobTimeout     = 24 * time.Hour
	defaultBidTimeout     = 5 * time.Second
	defaultTopK           = 3
	// re-poll in case a change notification is missed.
	placementPoll = 2 * time.Second
)

type Config struct {
	SharedSecret   string
	MaxPending     int
	ReconnectGrace time.Duration
	JobTimeout     time.Duration
	BidTimeout     time.Duration
	TopK           int
	LogDir         string
}

type Mill struct {
	l   *slog.Logger
	cfg Config

	// re-author relayed status rows into the mill's stream so appview sees
	// them unchanged.
	db *db.DB
	n  *notifier.Notifier

	mu         sync.Mutex
	sessions   map[string]*millSession
	leases     map[string]*RemoteLease
	nodeOffset map[string]uint64
	pending    int
	changeCh   chan struct{} // closed + replaced to wake placement waiters

	leaseSeq uint64
}

func New(l *slog.Logger, cfg Config) *Mill {
	if cfg.ReconnectGrace <= 0 {
		cfg.ReconnectGrace = defaultReconnectGrace
	}
	if cfg.JobTimeout <= 0 {
		cfg.JobTimeout = defaultJobTimeout
	}
	if cfg.BidTimeout <= 0 {
		cfg.BidTimeout = defaultBidTimeout
	}
	if cfg.TopK <= 0 {
		cfg.TopK = defaultTopK
	}
	return &Mill{
		l:          l,
		cfg:        cfg,
		sessions:   make(map[string]*millSession),
		leases:     make(map[string]*RemoteLease),
		nodeOffset: make(map[string]uint64),
		changeCh:   make(chan struct{}),
	}
}

func (m *Mill) Attach(d *db.DB, n *notifier.Notifier) {
	m.mu.Lock()
	m.db = d
	m.n = n
	m.mu.Unlock()
}

// dev bootstrap: shared-secret identity, single-token fleet. production
// should register per-executor tokens instead.
func (m *Mill) SeedBootstrapToken() error {
	if m.cfg.SharedSecret == "" || m.db == nil {
		return nil
	}
	return m.db.UpsertExecutorToken(bootstrapTokenName, HashToken(m.cfg.SharedSecret))
}

func (m *Mill) nextLeaseID() string {
	m.mu.Lock()
	m.leaseSeq++
	seq := m.leaseSeq
	m.mu.Unlock()
	return fmt.Sprintf("%s-%d", tid.TID(), seq)
}

func (m *Mill) notifyChange() {
	m.mu.Lock()
	m.notifyChangeLocked()
	m.mu.Unlock()
}

func (m *Mill) notifyChangeLocked() {
	close(m.changeCh)
	m.changeCh = make(chan struct{})
}

func (m *Mill) currentChangeCh() <-chan struct{} {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.changeCh
}

// --- session lifecycle -----------------------------------------------------

func (m *Mill) attachSession(sess *millSession) (uint64, bool) {
	m.mu.Lock()
	defer m.mu.Unlock()

	if old := m.sessions[sess.nodeID]; old != nil {
		if !old.disconnected {
			// reject a second live session for the same identity so a valid token
			// can't hijack an in-flight executor.
			return 0, false
		}
		// the old session is in its reconnect grace window: adopt its leases.
		if old.graceTimer != nil {
			old.graceTimer.Stop()
		}
		old.close()
		m.l.Info("executor reconnected", "node", sess.nodeID)
	}
	m.sessions[sess.nodeID] = sess
	return m.nodeOffset[sess.nodeID], true
}

func (m *Mill) detachSession(sess *millSession) {
	m.mu.Lock()
	if m.sessions[sess.nodeID] != sess {
		// already replaced by a reconnect; nothing to do.
		m.mu.Unlock()
		sess.close()
		return
	}
	sess.disconnected = true
	sess.graceTimer = time.AfterFunc(m.cfg.ReconnectGrace, func() { m.onGraceExpired(sess) })
	m.mu.Unlock()

	sess.close()
	m.l.Warn("executor session lost; entering reconnect grace", "node", sess.nodeID, "grace", m.cfg.ReconnectGrace)
	m.notifyChange()
}

func (m *Mill) sessionReady(sess *millSession) {
	for _, lease := range m.cancelledLeasesForNode(sess.nodeID) {
		_, reason := lease.cancelRequested()
		m.sendCancel(sess, lease, reason)
	}
	m.notifyChange()
}

func (m *Mill) cancelledLeasesForNode(nodeID string) []*RemoteLease {
	m.mu.Lock()
	var candidates []*RemoteLease
	for _, lease := range m.leases {
		if lease.nodeID == nodeID {
			candidates = append(candidates, lease)
		}
	}
	m.mu.Unlock()

	var leases []*RemoteLease
	for _, lease := range candidates {
		if cancelled, _ := lease.cancelRequested(); cancelled && lease.getState() != leaseDone {
			leases = append(leases, lease)
		}
	}
	return leases
}

func (m *Mill) onGraceExpired(sess *millSession) {
	m.mu.Lock()
	if m.sessions[sess.nodeID] != sess || !sess.disconnected {
		// reconnected in the meantime.
		m.mu.Unlock()
		return
	}
	var dead []*RemoteLease
	for _, lease := range m.leases {
		if lease.nodeID == sess.nodeID {
			dead = append(dead, lease)
		}
	}
	delete(m.sessions, sess.nodeID)
	delete(m.nodeOffset, sess.nodeID)
	m.mu.Unlock()

	m.l.Warn("executor declared dead; failing its in-flight jobs", "node", sess.nodeID, "jobs", len(dead))
	for _, lease := range dead {
		if cancelled, reason := lease.cancelRequested(); cancelled {
			lease.deliverCancelled(reason)
		} else {
			lease.markDead()
		}
	}
	m.notifyChange()
}

// --- placement (the WorkflowSlotter seam) ---------------------------------

func (m *Mill) place(ctx context.Context, engineName string, wid models.WorkflowId, wf *models.Workflow) (engine.WorkflowSlot, error) {
	m.mu.Lock()
	if m.cfg.MaxPending > 0 && m.pending >= m.cfg.MaxPending {
		max := m.cfg.MaxPending
		cur := m.pending
		m.mu.Unlock()
		return nil, fmt.Errorf("%w: mill has %d pending jobs (max %d)", engine.ErrNoWorkflowSlots, cur, max)
	}
	m.pending++
	m.mu.Unlock()
	defer func() {
		m.mu.Lock()
		m.pending--
		m.mu.Unlock()
	}()

	// stamp wid early so the executor reconstructs logs/status under the same
	// pipeline id the mill used for "pending".
	if st, ok := wf.Data.(*millWorkflowState); ok && st != nil {
		st.Wid = wid
	}

	for {
		if err := ctx.Err(); err != nil {
			return nil, err
		}

		lease, err := m.bid(ctx, engineName, wf)
		if err != nil {
			return nil, err
		}
		if lease != nil {
			lease.wid = wid
			m.mu.Lock()
			m.leases[lease.id] = lease
			if st, ok := wf.Data.(*millWorkflowState); ok && st != nil {
				st.Lease = lease
			}
			m.mu.Unlock()
			return &millSlot{fleet: m, lease: lease}, nil
		}

		// nobody to place on right now: wait for a change, ctx, or a poll tick.
		m.mu.Lock()
		ch := m.changeCh
		m.mu.Unlock()
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-ch:
		case <-time.After(placementPoll):
		}
	}
}

func (m *Mill) bid(ctx context.Context, engineName string, wf *models.Workflow) (*RemoteLease, error) {
	rawPipeline, rawWorkflow, err := marshalJob(wf)
	if err != nil {
		return nil, err
	}

	requiredLabels := requiredLabels(wf)
	requiredCapabilities := requiredCapabilities(wf)
	candidates := m.rankCandidates(engineName, requiredLabels, requiredCapabilities)
	if len(candidates) == 0 {
		return nil, nil
	}

	bidCtx, cancel := context.WithTimeout(ctx, m.cfg.BidTimeout)
	defer cancel()

	type bidResult struct {
		sess         *millSession
		lease        *RemoteLease
		rank         int
		incompatible bool
		reason       string
	}
	results := make(chan bidResult, len(candidates))
	ask := func(rank int, sess *millSession) {
		leaseID := m.nextLeaseID()
		lease := newLease(leaseID, sess.nodeID, engineName)
		wid := m.jobWid(wf)
		msg := &millproto.Message{ReserveSeat: &millv1.ReserveSeat{
			LeaseId:         leaseID,
			TargetEngine:    engineName,
			RawPipelineJson: rawPipeline,
			RawWorkflowJson: rawWorkflow,
			Knot:            wid.Knot,
			Rkey:            wid.Rkey,
			TtlSeconds:      uint32(m.cfg.ReconnectGrace / time.Second),
		}}
		resp, err := sess.request(bidCtx, leaseID, msg)
		if err != nil {
			results <- bidResult{}
			return
		}
		rr := resp.GetReserveResult()
		if rr == nil {
			results <- bidResult{}
			return
		}
		if !rr.GetAccepted() {
			if rr.GetRejectClass() == millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE {
				results <- bidResult{sess: sess, rank: rank, incompatible: true, reason: rr.GetRejectReason()}
				return
			}
			results <- bidResult{}
			return
		}
		results <- bidResult{sess: sess, lease: lease, rank: rank}
	}
	window := m.cfg.TopK
	if window <= 0 {
		window = len(candidates)
	}
	next := 0
	inFlight := 0
	for next < len(candidates) && inFlight < window {
		inFlight++
		go ask(next, candidates[next])
		next++
	}

	var winner *bidResult
	var losers []*RemoteLease
	var incompatible []string
	softReject := false
	for inFlight > 0 {
		r := <-results
		inFlight--
		if r.incompatible {
			if r.reason != "" {
				incompatible = append(incompatible, r.reason)
			}
		} else if r.lease == nil {
			softReject = true
		}
		if r.lease == nil {
			for winner == nil && next < len(candidates) && inFlight < window {
				inFlight++
				go ask(next, candidates[next])
				next++
			}
			continue
		}
		if winner == nil || r.rank < winner.rank {
			if winner != nil {
				losers = append(losers, winner.lease)
			}
			winner = &r
		} else {
			losers = append(losers, r.lease)
		}
	}

	// release the also-rans so they free their held seats immediately.
	for _, l := range losers {
		m.releaseRemote(l)
	}

	if winner == nil {
		if len(incompatible) > 0 && !softReject {
			return nil, fmt.Errorf("no compatible executor for %s: %s", engineName, strings.Join(incompatible, "; "))
		}
		return nil, nil
	}
	return winner.lease, nil
}

func (m *Mill) rankCandidates(engineName string, requiredLabels, requiredCapabilities []string) []*millSession {
	m.mu.Lock()
	defer m.mu.Unlock()

	type ranked struct {
		sess  *millSession
		worst float64
		sum   float64
	}
	var rs []ranked
	for _, s := range m.sessions {
		if s.disconnected {
			continue
		}
		if s.snapshot == nil {
			continue
		}
		ea, ok := s.snapshot.GetEngines()[engineName]
		if !ok || !ea.GetAvailable() {
			continue
		}
		if !hasLabels(s.labels, requiredLabels) {
			continue
		}
		if !hasLabels(ea.GetCapabilities(), requiredCapabilities) {
			continue
		}
		worst, sum := loadScore(ea.GetLoad())
		rs = append(rs, ranked{sess: s, worst: worst, sum: sum})
	}
	slices.SortStableFunc(rs, func(a, b ranked) int {
		if a.worst < b.worst {
			return -1
		}
		if a.worst > b.worst {
			return 1
		}
		if a.sum < b.sum {
			return -1
		}
		if a.sum > b.sum {
			return 1
		}
		return 0
	})

	out := make([]*millSession, len(rs))
	for i := range rs {
		out[i] = rs[i].sess
	}
	return out
}

func loadScore(load map[string]float64) (worst, sum float64) {
	for _, v := range load {
		if v > worst {
			worst = v
		}
		sum += v
	}
	return worst, sum
}

func requiredLabels(wf *models.Workflow) []string {
	st, ok := wf.Data.(*millWorkflowState)
	if !ok || st == nil {
		return nil
	}
	return st.RawWorkflow.RunsOn
}

func requiredCapabilities(wf *models.Workflow) []string {
	st, ok := wf.Data.(*millWorkflowState)
	if !ok || st == nil {
		return nil
	}
	return st.RequiredCapabilities
}

func hasLabels(labels []string, required []string) bool {
	for _, want := range required {
		if !slices.Contains(labels, want) {
			return false
		}
	}
	return true
}

func (m *Mill) commitAndWait(ctx context.Context, wf *models.Workflow, unlocked []secrets.UnlockedSecret) error {
	st, ok := wf.Data.(*millWorkflowState)
	if !ok || st == nil || st.Lease == nil {
		return fmt.Errorf("mill workflow state missing lease")
	}
	lease := st.Lease

	pbSecrets := make([]*millv1.Secret, len(unlocked))
	for i, s := range unlocked {
		pbSecrets[i] = &millv1.Secret{Key: s.Key, Value: s.Value}
	}

	commit := &millproto.Message{CommitLease: &millv1.CommitLease{
		LeaseId: lease.id,
		Secrets: pbSecrets,
	}}

	for {
		if res, ok := pollTerminal(lease); ok {
			return terminalError(res.GetTerminalStatus())
		}
		if !lease.markCommitting() {
			return engine.ErrWorkflowFailed
		}

		sess := m.sessionForNode(lease.nodeID)
		if sess == nil {
			if done, err := m.waitCommitRetry(ctx, lease); done || err != nil {
				return err
			}
			continue
		}

		reqCtx, cancel := context.WithTimeout(ctx, m.cfg.BidTimeout)
		resp, err := sess.request(reqCtx, lease.id, commit)
		cancel()
		if err != nil {
			switch {
			case errors.Is(err, errSessionClosed):
				if done, err := m.waitCommitRetry(ctx, lease); done || err != nil {
					return err
				}
				continue
			case errors.Is(err, context.DeadlineExceeded) && ctx.Err() == nil:
				continue
			case errors.Is(err, context.DeadlineExceeded):
				return engine.ErrTimedOut
			case errors.Is(err, context.Canceled):
				return err
			default:
				m.l.Warn("commit lease send failed; waiting for reconnect", "lease", lease.id, "node", lease.nodeID, "err", err)
				if done, err := m.waitCommitRetry(ctx, lease); done || err != nil {
					return err
				}
				continue
			}
		}
		if resp.GetCommitted() == nil {
			return engine.ErrWorkflowFailed
		}
		lease.markRunning()
		if cancelled, reason := lease.cancelRequested(); cancelled {
			m.sendCancel(sess, lease, reason)
		}
		break
	}

	select {
	case res := <-lease.terminal:
		return terminalError(res.GetTerminalStatus())
	case <-lease.dead:
		m.l.Warn("executor lost while running job", "lease", lease.id, "node", lease.nodeID)
		return engine.ErrWorkflowFailed
	case <-ctx.Done():
		if ctx.Err() == context.DeadlineExceeded {
			return engine.ErrTimedOut
		}
		return ctx.Err()
	}
}

func (m *Mill) waitCommitRetry(ctx context.Context, lease *RemoteLease) (bool, error) {
	res, err := m.waitLeaseChange(ctx, lease)
	if err != nil {
		return true, err
	}
	if res != nil {
		return true, terminalError(res.GetTerminalStatus())
	}
	return false, nil
}

func terminalError(status string) error {
	switch models.StatusKind(status) {
	case models.StatusKindSuccess:
		return nil
	case models.StatusKindTimeout:
		return engine.ErrTimedOut
	case models.StatusKindCancelled:
		return engine.ErrCancelled
	default:
		return engine.ErrWorkflowFailed
	}
}

func pollTerminal(lease *RemoteLease) (*millv1.AttemptResult, bool) {
	select {
	case res := <-lease.terminal:
		return res, true
	default:
		return nil, false
	}
}

func (m *Mill) waitLeaseChange(ctx context.Context, lease *RemoteLease) (*millv1.AttemptResult, error) {
	ch := m.currentChangeCh()
	select {
	case res := <-lease.terminal:
		return res, nil
	case <-lease.dead:
		return nil, engine.ErrWorkflowFailed
	case <-ctx.Done():
		if ctx.Err() == context.DeadlineExceeded {
			return nil, engine.ErrTimedOut
		}
		return nil, ctx.Err()
	case <-ch:
		return nil, nil
	case <-time.After(placementPoll):
		return nil, nil
	}
}

// --- cleanup ---------------------------------------------------------------

func (m *Mill) destroy(wid models.WorkflowId) {
	m.mu.Lock()
	var lease *RemoteLease
	for _, l := range m.leases {
		if l.wid == wid {
			lease = l
			break
		}
	}
	m.mu.Unlock()
	if lease == nil {
		return
	}
	reason := "workflow destroyed"
	switch lease.requestCancel(reason) {
	case cancelLocal:
		if sess := m.sessionForNode(lease.nodeID); sess != nil {
			_ = sess.send(&millproto.Message{ReleaseLease: &millv1.ReleaseLease{LeaseId: lease.id}})
		}
		lease.deliverCancelled(reason)
	case cancelRemote:
		if sess := m.sessionForNode(lease.nodeID); sess != nil {
			m.sendCancel(sess, lease, reason)
		}
	}
}

func (m *Mill) releaseSlot(s *millSlot) {
	lease := s.lease
	if lease.getState() == leaseReserved {
		m.releaseRemote(lease)
	}
	lease.closeLog()

	m.mu.Lock()
	delete(m.leases, lease.id)
	m.mu.Unlock()

	m.notifyChange()
}

func (m *Mill) releaseRemote(lease *RemoteLease) {
	lease.setState(leaseDone)
	if sess := m.sessionForNode(lease.nodeID); sess != nil {
		_ = sess.send(&millproto.Message{ReleaseLease: &millv1.ReleaseLease{LeaseId: lease.id}})
	}
}

func (m *Mill) sendCancel(sess *millSession, lease *RemoteLease, reason string) {
	_ = sess.send(&millproto.Message{CancelAttempt: &millv1.CancelAttempt{
		LeaseId: lease.id,
		Reason:  reason,
	}})
}

func (m *Mill) sessionForNode(nodeID string) *millSession {
	m.mu.Lock()
	defer m.mu.Unlock()
	s := m.sessions[nodeID]
	if s == nil || s.disconnected {
		return nil
	}
	return s
}

// --- relay handlers --------------------------------------------------------

func (m *Mill) shouldAcceptOffset(sess *millSession, offset uint64) bool {
	m.mu.Lock()
	defer m.mu.Unlock()
	current := m.nodeOffset[sess.nodeID]
	return offset == current+1
}

func (m *Mill) ackOffset(sess *millSession, offset uint64) {
	m.mu.Lock()
	if offset > m.nodeOffset[sess.nodeID] {
		m.nodeOffset[sess.nodeID] = offset
	}
	m.mu.Unlock()

	_ = sess.send(&millproto.Message{Ack: &millv1.Ack{UpToOffset: offset}})
}

func (m *Mill) processRelay(sess *millSession, offset uint64, kind string, apply func() error) {
	if !m.shouldAcceptOffset(sess, offset) {
		return
	}
	// ack only after the side effect lands, or a transient write error makes
	// the message unreplayable.
	if err := apply(); err != nil {
		m.l.Error("process relayed message failed", "kind", kind, "node", sess.nodeID, "offset", offset, "err", err)
		return
	}
	m.ackOffset(sess, offset)
}

func (m *Mill) onSnapshot(sess *millSession, snap *millv1.NodeSnapshot) {
	m.mu.Lock()
	sess.snapshot = snap
	m.mu.Unlock()
	m.notifyChange()
}

func (m *Mill) onStatusRelay(sess *millSession, ev *millv1.StatusEvent) {
	m.processRelay(sess, ev.GetOffset(), "status", func() error {
		if m.db == nil {
			return fmt.Errorf("mill db not attached")
		}
		m.mu.Lock()
		lease := m.leases[ev.GetLeaseId()]
		m.mu.Unlock()
		if lease == nil || lease.nodeID != sess.nodeID {
			// not a lease this executor holds; don't let it author status rows for
			// arbitrary pipelines.
			return nil
		}
		var errMsg *string
		if e := ev.GetError(); e != "" {
			errMsg = &e
		}
		var exit *int64
		if c := ev.GetExitCode(); c != 0 {
			exit = &c
		}
		// authoritative pipeline/workflow come from the lease, not whatever the
		// executor sent.
		return m.db.InsertRelayedStatus(string(lease.wid.PipelineId.AtUri()), lease.wid.Name, ev.GetStatus(), errMsg, exit, m.n)
	})
}

func (m *Mill) onLogRelay(sess *millSession, ll *millv1.LogLine) {
	m.processRelay(sess, ll.GetOffset(), "log", func() error {
		m.mu.Lock()
		lease := m.leases[ll.GetLeaseId()]
		m.mu.Unlock()
		if lease == nil || lease.nodeID != sess.nodeID {
			return nil
		}
		path := models.LogFilePath(m.cfg.LogDir, lease.wid)
		line := append(ll.GetRawJson(), '\n')
		return lease.writeLog(path, line)
	})
}

func (m *Mill) onAttemptResult(sess *millSession, ar *millv1.AttemptResult) {
	m.processRelay(sess, ar.GetOffset(), "attempt-result", func() error {
		m.mu.Lock()
		lease := m.leases[ar.GetLeaseId()]
		m.mu.Unlock()
		if lease == nil || lease.nodeID != sess.nodeID {
			return nil
		}
		lease.deliverTerminal(ar)
		return nil
	})
}

// --- helpers ---------------------------------------------------------------

func marshalJob(wf *models.Workflow) (pipeline string, workflow string, err error) {
	st, ok := wf.Data.(*millWorkflowState)
	if !ok || st == nil {
		return "", "", fmt.Errorf("mill workflow state missing")
	}
	p, err := json.Marshal(st.RawPipeline)
	if err != nil {
		return "", "", fmt.Errorf("marshal pipeline: %w", err)
	}
	w, err := json.Marshal(st.RawWorkflow)
	if err != nil {
		return "", "", fmt.Errorf("marshal workflow: %w", err)
	}
	return string(p), string(w), nil
}

func (m *Mill) jobWid(wf *models.Workflow) models.WorkflowId {
	st, ok := wf.Data.(*millWorkflowState)
	if !ok || st == nil {
		return models.WorkflowId{}
	}
	return st.Wid
}
