package executor

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"maps"
	"net/http"
	"runtime"
	"strings"
	"sync"
	"time"

	"github.com/bluesky-social/indigo/atproto/syntax"
	"github.com/gorilla/websocket"

	"tangled.org/core/api/tangled"
	"tangled.org/core/notifier"
	"tangled.org/core/spindle/config"
	"tangled.org/core/spindle/db"
	"tangled.org/core/spindle/engine"
	"tangled.org/core/spindle/models"

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

const (
	dialBackoffMin = 1 * time.Second
	dialBackoffMax = 30 * time.Second
	snapshotEvery  = 15 * time.Second
	defaultSeats   = 4
)

type Executor struct {
	millURL string
	token   string
	nodeID  string
	seats   int
	labels  []string

	engines map[string]models.Engine
	db      *db.DB
	n       *notifier.Notifier
	cfg     *config.Config
	l       *slog.Logger

	relay   *relayLog
	relayMu sync.Mutex // offset assign + send must be serialized so wire order == offset order

	connMu sync.Mutex
	enc    messageEncoder

	mu       sync.Mutex
	active   map[string]*reservation
	draining bool
}

type reservation struct {
	leaseID    string
	wid        models.WorkflowId
	realEngine models.Engine
	slot       engine.WorkflowSlot
	wf         *models.Workflow
	repoDid    syntax.DID

	committed bool
	cancelled bool
	cancel    context.CancelFunc
	ttlTimer  *time.Timer
	stopTail  func()
}

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

func New(cfg *config.Config, engines map[string]models.Engine, d *db.DB, n *notifier.Notifier, l *slog.Logger) *Executor {
	seats := cfg.Mill.Seats
	if seats <= 0 {
		seats = defaultSeats
	}
	labels := normalizeLabels(cfg.Mill.Labels)
	return &Executor{
		millURL: cfg.Mill.URL,
		token:   cfg.Mill.SharedSecret,
		nodeID:  cfg.Server.Hostname,
		seats:   seats,
		labels:  labels,
		engines: engines,
		db:      d,
		n:       n,
		cfg:     cfg,
		l:       l.With("component", "mill.executor"),
		relay:   newRelayLog(),
		active:  make(map[string]*reservation),
	}
}

func (e *Executor) Connect(ctx context.Context) {
	go e.observeLoop(ctx)

	backoff := dialBackoffMin
	for {
		if ctx.Err() != nil {
			return
		}
		err := e.runSession(ctx)
		if ctx.Err() != nil {
			return
		}
		e.l.Warn("mill session ended; reconnecting", "err", err, "backoff", backoff)
		select {
		case <-ctx.Done():
			return
		case <-time.After(backoff):
		}
		backoff = min(backoff*2, dialBackoffMax)
	}
}

func (e *Executor) runSession(ctx context.Context) error {
	header := http.Header{}
	if e.token != "" {
		header.Set("Authorization", "Bearer "+e.token)
	}
	conn, _, err := websocket.DefaultDialer.DialContext(ctx, e.millURL, header)
	if err != nil {
		return fmt.Errorf("dial mill: %w", err)
	}
	defer conn.Close()

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

	hello := &millproto.Message{Hello: &millv1.Hello{
		ProtocolVersion: millproto.ProtocolVersion,
		Arch:            runtime.GOARCH,
		Labels:          e.labels,
	}}
	if err := enc.Encode(hello); err != nil {
		return fmt.Errorf("send hello: %w", err)
	}

	resumeMsg, err := dec.Decode()
	if err != nil {
		return fmt.Errorf("read resume: %w", err)
	}
	resume := resumeMsg.GetResume()
	if resume == nil {
		return fmt.Errorf("expected resume, got something else")
	}

	// install encoder and replay anything the mill missed.
	e.connMu.Lock()
	e.enc = enc
	e.connMu.Unlock()
	defer func() {
		e.connMu.Lock()
		e.enc = nil
		e.connMu.Unlock()
	}()

	e.replay(resume.GetAckOffset())
	e.pushSnapshot()
	e.l.Info("connected to mill", "node", e.nodeID, "resumeFrom", resume.GetAckOffset())

	go e.snapshotLoop(ctx, enc)

	for {
		msg, err := dec.Decode()
		if err != nil {
			return fmt.Errorf("read: %w", err)
		}
		e.dispatch(ctx, msg)
	}
}

func (e *Executor) replay(ackOffset uint64) {
	e.relayMu.Lock()
	defer e.relayMu.Unlock()
	e.relay.ack(ackOffset)
	for _, msg := range e.relay.since(ackOffset) {
		e.send(msg)
	}
}

func (e *Executor) send(msg *millproto.Message) {
	e.connMu.Lock()
	enc := e.enc
	e.connMu.Unlock()
	if enc != nil {
		_ = enc.Encode(msg)
	}
}

func (e *Executor) dispatch(ctx context.Context, msg *millproto.Message) {
	switch {
	case msg.GetReserveSeat() != nil:
		e.handleReserve(ctx, msg.GetReserveSeat())
	case msg.GetCommitLease() != nil:
		e.handleCommit(ctx, msg.GetCommitLease())
	case msg.GetReleaseLease() != nil:
		e.handleRelease(msg.GetReleaseLease().GetLeaseId())
	case msg.GetCancelAttempt() != nil:
		e.handleCancel(msg.GetCancelAttempt().GetLeaseId())
	case msg.GetAck() != nil:
		e.relay.ack(msg.GetAck().GetUpToOffset())
	default:
		e.l.Warn("executor received unexpected message")
	}
}

// --- reserve / commit / release / cancel ----------------------------------

func (e *Executor) handleReserve(ctx context.Context, rs *millv1.ReserveSeat) {
	reject := func(reason string, class millv1.RejectClass) {
		e.send(&millproto.Message{ReserveResult: &millv1.ReserveResult{
			LeaseId:      rs.GetLeaseId(),
			Accepted:     false,
			RejectReason: reason,
			RejectClass:  class,
		}})
	}

	e.mu.Lock()
	draining := e.draining
	e.mu.Unlock()
	if draining {
		reject("draining", millv1.RejectClass_REJECT_CLASS_TRANSIENT)
		return
	}

	realEngine, ok := e.engines[rs.GetTargetEngine()]
	if !ok {
		reject("unknown engine "+rs.GetTargetEngine(), millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE)
		return
	}
	slotter, ok := realEngine.(engine.WorkflowSlotter)
	if !ok {
		reject("engine does not support workflow slots", millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE)
		return
	}

	var twf tangled.Pipeline_Workflow
	if err := json.Unmarshal([]byte(rs.GetRawWorkflowJson()), &twf); err != nil {
		reject("bad workflow json", millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE)
		return
	}
	var tpl tangled.Pipeline
	if err := json.Unmarshal([]byte(rs.GetRawPipelineJson()), &tpl); err != nil {
		reject("bad pipeline json", millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE)
		return
	}

	pipelineId := models.PipelineId{Knot: rs.GetKnot(), Rkey: rs.GetRkey()}
	wid := models.WorkflowId{PipelineId: pipelineId, Name: twf.Name}

	wf, err := realEngine.InitWorkflow(twf, tpl)
	if err != nil {
		reject("init workflow: "+err.Error(), millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE)
		return
	}
	if validator, ok := realEngine.(engine.WorkflowPlacementValidator); ok {
		if err := validator.ValidateWorkflowPlacement(wf); err != nil {
			reject("validate workflow placement: "+err.Error(), millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE)
			return
		}
	}
	// the job skipped processPipeline, so inject TANGLED_* env here.
	if wf.Environment == nil {
		wf.Environment = make(map[string]string)
	}
	maps.Copy(wf.Environment, models.PipelineEnvVars(tpl.TriggerMetadata, pipelineId))

	// NoWait: the executor doesn't queue locally; the mill owns the backlog.
	slot, err := slotter.AcquireWorkflowSlot(ctx, wid, wf, engine.NoWait)
	if err != nil {
		class := millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE
		if errors.Is(err, engine.ErrNoWorkflowSlots) {
			class = millv1.RejectClass_REJECT_CLASS_TRANSIENT
		}
		reject(err.Error(), class)
		return
	}

	var repoDid syntax.DID
	if tpl.TriggerMetadata != nil && tpl.TriggerMetadata.Repo != nil && tpl.TriggerMetadata.Repo.RepoDid != nil {
		repoDid, _ = syntax.ParseDID(*tpl.TriggerMetadata.Repo.RepoDid)
	}

	res := &reservation{
		leaseID:    rs.GetLeaseId(),
		wid:        wid,
		realEngine: realEngine,
		slot:       slot,
		wf:         wf,
		repoDid:    repoDid,
	}
	res.ttlTimer = time.AfterFunc(ttlDuration(rs.GetTtlSeconds()), func() { e.expireReservation(res.leaseID) })

	e.mu.Lock()
	e.active[res.leaseID] = res
	e.mu.Unlock()

	e.send(&millproto.Message{ReserveResult: &millv1.ReserveResult{
		LeaseId:  rs.GetLeaseId(),
		Accepted: true,
	}})
	e.pushSnapshot()
}

func (e *Executor) handleCommit(ctx context.Context, cl *millv1.CommitLease) {
	e.mu.Lock()
	res := e.active[cl.GetLeaseId()]
	if res == nil {
		e.mu.Unlock()
		return
	}
	if res.committed {
		e.mu.Unlock()
		e.send(&millproto.Message{Committed: &millv1.Committed{LeaseId: cl.GetLeaseId()}})
		return
	}
	res.committed = true
	if res.ttlTimer != nil {
		res.ttlTimer.Stop()
	}
	jobCtx, cancel := context.WithCancel(ctx)
	res.cancel = cancel
	e.mu.Unlock()

	vault := newMemVault(cl.GetSecrets())
	re := newReservedEngine(res.realEngine, res.slot)
	pipeline := &models.Pipeline{
		RepoDid:   res.repoDid,
		Workflows: map[models.Engine][]models.Workflow{re: {*res.wf}},
	}

	e.startTail(res)

	go engine.StartWorkflows(e.l, vault, e.cfg, e.db, e.n, jobCtx, pipeline, res.wid.PipelineId)

	e.send(&millproto.Message{Committed: &millv1.Committed{LeaseId: cl.GetLeaseId()}})
}

func (e *Executor) handleRelease(leaseID string) {
	cleanup, ok := e.takeUncommittedReservation(leaseID, true)
	if !ok {
		return
	}
	cleanup()
	e.pushSnapshot()
}

func (e *Executor) handleCancel(leaseID string) {
	e.mu.Lock()
	res := e.active[leaseID]
	if res == nil {
		e.mu.Unlock()
		return
	}
	res.cancelled = true
	cancel := res.cancel
	committed := res.committed
	var cleanup func()
	if !committed {
		cleanup = e.removeReservationLocked(res, true)
	}
	e.mu.Unlock()

	if cancel != nil {
		cancel()
	}
	// committed jobs clean up when the observe loop sees the terminal row; an
	// uncommitted cancel still needs the slot released.
	if !committed {
		cleanup()
		e.pushSnapshot()
	}
}

func (e *Executor) expireReservation(leaseID string) {
	cleanup, ok := e.takeUncommittedReservation(leaseID, true)
	if !ok {
		return
	}
	e.l.Warn("reservation expired before commit", "lease", leaseID)
	cleanup()
	e.pushSnapshot()
}

func (e *Executor) takeUncommittedReservation(leaseID string, releaseSlot bool) (func(), bool) {
	e.mu.Lock()
	defer e.mu.Unlock()
	res := e.active[leaseID]
	if res == nil || res.committed {
		return func() {}, false
	}
	return e.removeReservationLocked(res, releaseSlot), true
}

func (e *Executor) finishReservation(res *reservation) (func(), bool, bool) {
	e.mu.Lock()
	defer e.mu.Unlock()
	if e.active[res.leaseID] != res {
		return func() {}, false, false
	}
	cancelled := res.cancelled
	return e.removeReservationLocked(res, false), cancelled, true
}

func (e *Executor) removeReservationLocked(res *reservation, releaseSlot bool) func() {
	delete(e.active, res.leaseID)
	if res.ttlTimer != nil {
		res.ttlTimer.Stop()
		res.ttlTimer = nil
	}
	stopTail := res.stopTail
	res.stopTail = nil
	slot := res.slot
	if releaseSlot {
		res.slot = nil
	}
	return func() {
		if stopTail != nil {
			stopTail()
		}
		if releaseSlot && slot != nil {
			slot.Release()
		}
	}
}

func normalizeLabels(labels []string) []string {
	seen := make(map[string]struct{}, len(labels))
	out := make([]string, 0, len(labels))
	for _, label := range labels {
		label = strings.TrimSpace(label)
		if label == "" {
			continue
		}
		if _, ok := seen[label]; ok {
			continue
		}
		seen[label] = struct{}{}
		out = append(out, label)
	}
	return out
}

// --- snapshots -------------------------------------------------------------

func (e *Executor) snapshotLoop(ctx context.Context, enc *millproto.Encoder) {
	t := time.NewTicker(snapshotEvery)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			// stop pushing once this connection is replaced.
			e.connMu.Lock()
			cur := e.enc
			e.connMu.Unlock()
			if cur != enc {
				return
			}
			e.pushSnapshot()
		}
	}
}

func (e *Executor) pushSnapshot() {
	e.mu.Lock()
	draining := e.draining
	active := len(e.active)
	e.mu.Unlock()

	load := 0.0
	if e.seats > 0 {
		load = float64(active) / float64(e.seats)
	}
	available := !draining && active < e.seats
	engines := make(map[string]*millv1.EngineAvailability, len(e.engines))
	for name, realEngine := range e.engines {
		var capabilities []string
		if provider, ok := realEngine.(engine.CapabilityProvider); ok {
			var err error
			capabilities, err = engine.NormalizeCapabilities(provider.Capabilities())
			if err != nil {
				e.l.Error("engine advertised invalid placement capabilities", "engine", name, "err", err)
				capabilities = nil
			}
		}
		engines[name] = &millv1.EngineAvailability{
			Available:    available,
			Load:         map[string]float64{"slots": load},
			Capabilities: capabilities,
		}
	}
	e.send(&millproto.Message{NodeSnapshot: &millv1.NodeSnapshot{
		NodeId:  e.nodeID,
		Engines: engines,
	}})
}

func (e *Executor) Drain() {
	e.mu.Lock()
	e.draining = true
	e.mu.Unlock()
	e.pushSnapshot()
}

func ttlDuration(secs uint32) time.Duration {
	if secs == 0 {
		return defaultReservationTTL
	}
	return time.Duration(secs) * time.Second
}

const defaultReservationTTL = 60 * time.Second
