package orbit

import (
	"context"
	"errors"
	"math/rand/v2"
	"sync"
	"sync/atomic"
	"time"
)

type runtimeAttempt struct {
	assignment                 WorkerAssignment
	startID, renewID, finishID string
	guard                      *executionGuard
	invoked, returned, lost    bool
	result                     WorkerFinish
	nextTry, nextRenew         time.Time
	retryDelay                 time.Duration
}
type runtimeCommand struct {
	kind     string
	id       string
	job      *runtimeAttempt
	capacity WorkerCapacity
	conn     *WorkerConnection
}
type runtimeReply struct {
	command *runtimeCommand
	result  WorkerCommandResult
	welcome WorkerWelcome
	err     error
}
type runtimeConnected struct {
	conn   *WorkerConnection
	leases map[string]WorkerLease
	err    error
}

type workerRuntime struct {
	worker                   *Worker
	ctx                      context.Context
	jobs                     map[string]*runtimeAttempt
	order                    []string
	cursor                   int
	conn                     *WorkerConnection
	command                  *runtimeCommand
	connecting               bool
	connectedOnce            bool
	draining, drainConfirmed bool
	capacity                 uint32
	sequence                 uint64
	heartbeat, dialAt        time.Time
	backoff                  time.Duration
	results                  chan handlerCompletion
	replies                  chan runtimeReply
	connections              chan runtimeConnected
	wg                       sync.WaitGroup
}

// Run 持续消费直到 ctx 取消或不可恢复的认证/协议错误。取消后先排空并在
// ShutdownTimeout 内收集在途结果；正常收尾返回 nil，超时返回 *WorkerRunError。
// 已到期的处理器上下文不会因重连复活；每个已知尝试在本进程最多调用一次处理器。
func (w *Worker) Run(parent context.Context) error {
	if !w.used.CompareAndSwap(false, true) {
		return ErrWorkerAlreadyRun
	}
	defer close(w.events)
	if parent.Err() != nil {
		return nil
	}
	ctx, cancel := context.WithCancel(context.WithoutCancel(parent))
	var shutdownExpired atomic.Bool
	watchDone := make(chan struct{})
	go func() {
		defer close(watchDone)
		select {
		case <-ctx.Done():
			return
		case <-parent.Done():
		}
		timer := time.NewTimer(w.options.ShutdownTimeout)
		defer timer.Stop()
		select {
		case <-ctx.Done():
		case <-timer.C:
			shutdownExpired.Store(true)
			cancel()
		}
	}()
	r := &workerRuntime{worker: w, ctx: ctx, jobs: make(map[string]*runtimeAttempt), backoff: w.options.ReconnectMinDelay, results: make(chan handlerCompletion, w.options.Capacity), replies: make(chan runtimeReply, 2), connections: make(chan runtimeConnected, 1)}
	defer func() {
		cancel()
		if r.conn != nil {
			r.conn.Close()
		}
		for _, j := range r.jobs {
			if j.guard != nil {
				j.guard.stop(context.Canceled)
			}
		}
		r.wg.Wait()
		// 连接可能在退出与交付结果之间成功，仍需关闭；不留下新会话连接。
		select {
		case c := <-r.connections:
			if c.conn != nil {
				c.conn.Close()
			}
		default:
		}
		<-watchDone
	}()
	tick := time.NewTicker(20 * time.Millisecond)
	defer tick.Stop()
	stop := parent.Done()
	for {
		if ctx.Err() != nil {
			cause := ctx.Err()
			if shutdownExpired.Load() {
				cause = ErrWorkerShutdownTimeout
			}
			return r.failure(cause)
		}
		if parent.Err() != nil && !r.draining {
			r.draining = true
			stop = nil
			w.emit(WorkerEvent{Type: "draining"})
		}
		r.observeExpired()
		if r.conn != nil && r.conn.Err() != nil {
			if err := r.disconnect(r.conn.Err()); err != nil {
				return r.failure(err)
			}
		}
		// drain ACK 之前的派发已被底层读取器排进队列，收齐后才能宣布排空。
		if r.conn != nil {
			for len(r.conn.assignments) > 0 {
				if err := r.accept(<-r.conn.assignments); err != nil {
					return r.failure(err)
				}
			}
		}
		if r.draining && len(r.jobs) == 0 && (r.drainConfirmed || !r.connectedOnce && !r.connecting) {
			return nil
		}
		if r.conn == nil && !r.connecting && !time.Now().Before(r.dialAt) {
			r.connect()
		}
		if r.conn != nil && r.command == nil {
			if err := r.dispatch(); err != nil {
				return r.failure(err)
			}
		}
		var assignments <-chan WorkerAssignment
		var closed <-chan struct{}
		if r.conn != nil {
			assignments = r.conn.assignments
			closed = r.conn.Done()
		}
		select {
		case <-ctx.Done():
		case <-stop:
			stop = nil
		case <-tick.C:
		case <-closed:
		case a := <-assignments:
			if r.conn.Err() == nil {
				if err := r.accept(a); err != nil {
					return r.failure(err)
				}
			}
		case done := <-r.results:
			if j := r.jobs[done.attemptID]; j != nil {
				j.returned = true
				j.result = done.result
				j.nextTry = time.Time{}
			}
		case reply := <-r.replies:
			if err := r.receive(reply); err != nil {
				return r.failure(err)
			}
		case connected := <-r.connections:
			if err := r.establish(connected); err != nil {
				return r.failure(err)
			}
		}
	}
}

func (r *workerRuntime) failure(err error) error {
	return &WorkerRunError{Cause: err, Unconfirmed: len(r.jobs), Running: r.worker.running.Load()}
}
func (r *workerRuntime) event(j *runtimeAttempt, kind, code string) {
	r.worker.emit(WorkerEvent{Type: kind, TaskID: j.assignment.Task.ID, AttemptID: j.assignment.AttemptID, Code: code})
}
func (r *workerRuntime) accept(a WorkerAssignment) error {
	if j := r.jobs[a.AttemptID]; j != nil {
		if j.assignment.Attempt() != a.Attempt() {
			return ErrWorkerProtocol
		}
		return nil
	}
	// 运行中失去授权但不退出的处理器仍占物理槽；允许接收一轮已在网络中的预占，
	// 不突破处理器容量，也不让待核对身份无限累积。
	if len(r.jobs) >= 2*int(r.worker.options.Capacity) {
		return ErrWorkerBackpressure
	}
	id, err := NewWorkerMessageID()
	if err != nil {
		return err
	}
	r.jobs[a.AttemptID] = &runtimeAttempt{assignment: a, startID: id}
	r.order = append(r.order, a.AttemptID)
	return nil
}
func (r *workerRuntime) remove(j *runtimeAttempt) {
	if j.guard != nil {
		j.guard.stop(context.Canceled)
	}
	delete(r.jobs, j.assignment.AttemptID)
	for i, id := range r.order {
		if id == j.assignment.AttemptID {
			r.order = append(r.order[:i], r.order[i+1:]...)
			break
		}
	}
	if len(r.order) > 0 {
		r.cursor %= len(r.order)
	} else {
		r.cursor = 0
	}
}
func (r *workerRuntime) observeExpired() {
	for _, j := range r.jobs {
		if j.guard != nil && j.guard.ctx.Err() != nil && !j.lost {
			j.lost = true
			r.event(j, "lease_lost", "LEASE_LOST")
		}
	}
}
func (r *workerRuntime) desiredCapacity() uint32 {
	if r.draining {
		return 0
	}
	ghosts := uint32(0)
	for _, j := range r.jobs {
		if j.returned || j.lost {
			ghosts++
		}
	}
	if ghosts >= r.worker.options.Capacity {
		return 0
	}
	return r.worker.options.Capacity - ghosts
}

func (r *workerRuntime) connect() {
	r.connecting = true
	if r.connectedOnce {
		r.worker.reconnects.Add(1)
		r.worker.emit(WorkerEvent{Type: "reconnecting"})
	}
	refs := make([]WorkerAttempt, 0, len(r.jobs))
	for _, j := range r.jobs {
		refs = append(refs, j.assignment.Attempt())
	}
	r.wg.Add(1)
	go func() {
		defer r.wg.Done()
		ctx, stop := context.WithTimeout(r.ctx, 30*time.Second)
		defer stop()
		c, err := r.worker.client.DialWorker(ctx, r.worker.registration)
		result := runtimeConnected{conn: c, err: err, leases: make(map[string]WorkerLease, len(refs))}
		if err == nil {
			for _, ref := range refs {
				id, e := NewWorkerMessageID()
				if e != nil {
					result.err = e
					break
				}
				lease, e := c.Resume(ctx, id, ref)
				if e != nil {
					result.err = e
					break
				}
				result.leases[ref.AttemptID] = lease
			}
		}
		if result.err != nil && c != nil {
			c.Close()
			result.conn = nil
		}
		select {
		case r.connections <- result:
		case <-r.ctx.Done():
			if result.conn != nil {
				result.conn.Close()
			}
		}
	}()
}
func (r *workerRuntime) establish(c runtimeConnected) error {
	r.connecting = false
	if c.err != nil {
		if !workerRetryable(c.err) {
			return c.err
		}
		r.deferReconnect()
		return nil
	}
	r.conn = c.conn
	r.connectedOnce = true
	r.sequence = 0
	r.heartbeat = time.Time{}
	r.capacity = ^uint32(0)
	r.drainConfirmed = false
	r.draining = r.draining || c.conn.Welcome().Draining
	for id, lease := range c.leases {
		j := r.jobs[id]
		if j == nil {
			continue
		}
		if j.guard != nil {
			if !j.guard.update(lease) {
				j.lost = true
				r.event(j, "lease_lost", "LEASE_LOST")
			}
			j.nextRenew = time.Time{}
		}
		if !j.invoked && !lease.Active {
			r.event(j, "lease_lost", "NOT_STARTED")
			r.remove(j)
		}
	}
	r.worker.emit(WorkerEvent{Type: "connected"})
	return nil
}
func (r *workerRuntime) deferReconnect() {
	min := r.worker.options.ReconnectMinDelay
	r.dialAt = time.Now().Add(min + time.Duration(rand.Int64N(int64(r.backoff-min)+1)))
	r.backoff = minDuration(r.backoff*2, r.worker.options.ReconnectMaxDelay)
}
func (r *workerRuntime) disconnect(err error) error {
	if r.conn != nil {
		r.conn.Close()
		r.conn = nil
	}
	r.command = nil
	if !workerRetryable(err) {
		return err
	}
	r.deferReconnect()
	return nil
}
func workerRetryable(err error) bool {
	if errors.Is(err, ErrWorkerDisconnected) || errors.Is(err, ErrWorkerClosed) || errors.Is(err, ErrWorkerBackpressure) || errors.Is(err, context.DeadlineExceeded) {
		return true
	}
	var control *WorkerError
	if errors.As(err, &control) {
		return control.Code == "STALE_SESSION" || control.Code == "UNAVAILABLE"
	}
	var httpErr *HTTPError
	return errors.As(err, &httpErr) && (httpErr.StatusCode == 429 || httpErr.StatusCode >= 500)
}
func minDuration(a, b time.Duration) time.Duration {
	if a < b {
		return a
	}
	return b
}

func (r *workerRuntime) dispatch() error {
	now := time.Now()
	capacity := r.desiredCapacity()
	if !now.Before(r.heartbeat) || capacity != r.capacity || r.draining && !r.drainConfirmed {
		id, err := NewWorkerMessageID()
		if err != nil {
			return err
		}
		r.sequence++
		r.launch(&runtimeCommand{kind: "capacity", id: id, capacity: WorkerCapacity{Sequence: r.sequence, Capacity: capacity, Drain: r.draining}})
		return nil
	}
	// 结果优先释放占用；其后照顾临近续租，再批准新处理器。每类内部轮转。
	for _, kind := range []string{"finish", "renew", "start"} {
		for n := 0; n < len(r.order); n++ {
			i := (r.cursor + n) % len(r.order)
			j := r.jobs[r.order[i]]
			if now.Before(j.nextTry) {
				continue
			}
			var id *string
			switch kind {
			case "finish":
				if !j.returned {
					continue
				}
				id = &j.finishID
			case "renew":
				if !j.invoked || j.lost || j.returned || now.Before(j.nextRenew) {
					continue
				}
				id = &j.renewID
			case "start":
				if j.invoked || r.worker.running.Load() >= int64(r.worker.options.Capacity) {
					continue
				}
				id = &j.startID
			}
			if *id == "" {
				var err error
				*id, err = NewWorkerMessageID()
				if err != nil {
					return err
				}
			}
			r.cursor = (i + 1) % len(r.order)
			r.launch(&runtimeCommand{kind: kind, id: *id, job: j})
			return nil
		}
	}
	return nil
}
func (r *workerRuntime) launch(command *runtimeCommand) {
	command.conn = r.conn
	r.command = command
	// 操作所需值在 goroutine 启动前复制；运行器事件循环独占可变尝试状态。
	var ref WorkerAttempt
	var finish WorkerFinish
	if command.job != nil {
		ref = command.job.assignment.Attempt()
		finish = command.job.result
	}
	r.wg.Add(1)
	go func() {
		defer r.wg.Done()
		reply := runtimeReply{command: command}
		switch command.kind {
		case "capacity":
			reply.welcome, reply.err = command.conn.Capacity(r.ctx, command.id, command.capacity)
		case "start":
			reply.result, reply.err = command.conn.Start(r.ctx, command.id, ref)
		case "renew":
			reply.result, reply.err = command.conn.Renew(r.ctx, command.id, ref)
		case "finish":
			reply.result, reply.err = command.conn.Finish(r.ctx, command.id, finish)
		}
		select {
		case r.replies <- reply:
		case <-r.ctx.Done():
		}
	}()
}

func (r *workerRuntime) receive(reply runtimeReply) error {
	command := reply.command
	if r.command != command || command.conn != r.conn {
		return nil
	}
	r.command = nil
	j := command.job
	if reply.err != nil {
		var control *WorkerError
		if j != nil && errors.As(reply.err, &control) {
			switch control.Code {
			case "LEASE_LOST":
				if command.kind == "finish" {
					return reply.err
				}
				if j.guard != nil {
					j.guard.stop(ErrExecutionLeaseLost)
					j.lost = true
					r.event(j, "lease_lost", control.Code)
				} else {
					r.remove(j)
				}
				return nil
			case "EXECUTION_PAUSED", "UNAVAILABLE":
				if j.retryDelay == 0 {
					j.retryDelay = 200 * time.Millisecond
				} else {
					j.retryDelay = minDuration(j.retryDelay*2, 5*time.Second)
				}
				j.nextTry = time.Now().Add(j.retryDelay)
				return nil
			}
		}
		if workerRetryable(reply.err) {
			return r.disconnect(reply.err)
		}
		return reply.err
	}
	r.backoff = r.worker.options.ReconnectMinDelay
	if command.kind == "capacity" {
		r.capacity = command.capacity.Capacity
		r.drainConfirmed = command.capacity.Drain && reply.welcome.Draining
		r.heartbeat = time.Now().Add(minDuration(5*time.Second, reply.welcome.ExpiresAt.Sub(reply.welcome.ServerTime)/3))
		return nil
	}
	if command.kind == "finish" {
		if reply.result.Receipt.Disposition == "APPLIED" {
			r.worker.reports.Add(1)
			r.event(j, "result_applied", string(reply.result.Receipt.TaskState))
		} else {
			r.worker.evidence.Add(1)
			r.event(j, "result_evidence", string(reply.result.Receipt.TaskState))
		}
		r.remove(j)
		return nil
	}
	if command.kind == "start" {
		guard, ok := newExecutionGuard(r.ctx, reply.result.Lease)
		if !ok {
			r.event(j, "lease_lost", "START_NOT_AUTHORIZED")
			r.remove(j)
			return nil
		}
		j.guard = guard
		j.invoked = true
		r.worker.running.Add(1)
		task := j.assignment.Task
		task.Payload = append([]byte(nil), task.Payload...)
		task.Result = append([]byte(nil), task.Result...)
		execution := WorkerExecution{Task: task, AttemptID: j.assignment.AttemptID, Epoch: j.assignment.Epoch}
		handler := r.worker.handlers[WorkerCapability{Handler: task.Handler, HandlerVersion: task.HandlerVersion, PayloadVersion: task.PayloadVersion}]
		r.event(j, "handler_started", "")
		go r.worker.invoke(guard.ctx, handler, execution, j.assignment.Attempt(), r.results, r.ctx.Done())
	} else {
		r.worker.renewals.Add(1)
		j.renewID = ""
		if !j.guard.update(reply.result.Lease) {
			j.lost = true
			r.event(j, "lease_lost", "LEASE_LOST")
		}
	}
	j.nextTry = time.Time{}
	j.retryDelay = 0
	if !j.lost {
		j.nextRenew = time.Now().Add(max(20*time.Millisecond, time.Until(j.guard.until())/3))
	}
	return nil
}