package node

import (
	"context"
	"errors"
	"sync"
	"time"

	"go.uber.org/zap"

	"iop/packages/go/credentiallease"
	runtime "iop/packages/go/execution"
	iop "iop/proto/gen/iop"
)

// defaultAttemptCloseGrace bounds the wait after cancellation used to decide
// whether an adapter has actually relinquished local ownership.
const defaultAttemptCloseGrace = 5 * time.Second

var errProviderResponseStalled = errors.New("provider response stalled")

// attemptClock is intentionally small so package tests can provide a manual
// timer without relying on scheduler sleeps.
type attemptClock interface {
	NewTimer(time.Duration) attemptTimer
	Now() time.Time
}
type attemptTimer interface {
	C() <-chan time.Time
	Stop() bool
	Reset(time.Duration) bool
}
type realAttemptClock struct{}
type realAttemptTimer struct{ timer *time.Timer }

func (realAttemptClock) NewTimer(d time.Duration) attemptTimer {
	return realAttemptTimer{timer: time.NewTimer(d)}
}
func (realAttemptClock) Now() time.Time               { return time.Now() }
func (t realAttemptTimer) C() <-chan time.Time        { return t.timer.C }
func (t realAttemptTimer) Stop() bool                 { return t.timer.Stop() }
func (t realAttemptTimer) Reset(d time.Duration) bool { return t.timer.Reset(d) }

// attemptObserver owns one response activity timer. Terminal and fenced
// states are monotonic so late provider output can never revive an attempt.
type attemptObserver struct {
	mu        sync.Mutex
	clock     attemptClock
	timer     attemptTimer
	deadline  time.Duration
	expiresAt time.Time
	terminal  bool
	fenced    bool
	epoch     uint64

	// beforeExpiryCapture and afterExpiryCapture are deterministic ordering
	// seams for package tests that exercise the receive-before-capture race.
	beforeExpiryCapture func()
	afterExpiryCapture  func(bool)
}

func newAttemptObserver(clock attemptClock, timeout time.Duration) *attemptObserver {
	if clock == nil {
		clock = realAttemptClock{}
	}
	// Record the arm's scheduled deadline before creating its timer. A very
	// short timer can signal while NewTimer is still returning; that signal is
	// nevertheless the current arm and must not be rejected as stale.
	armedAt := clock.Now()
	o := &attemptObserver{clock: clock, deadline: timeout, expiresAt: armedAt.Add(timeout)}
	o.timer = clock.NewTimer(timeout)
	return o
}
func (o *attemptObserver) expired() <-chan time.Time { return o.timer.C() }
func (o *attemptObserver) observe(disposition runtime.ProviderActivityDisposition) {
	o.mu.Lock()
	defer o.mu.Unlock()
	if o.terminal || o.fenced {
		return
	}
	switch disposition {
	case runtime.DispositionProgress:
		// A timer can have an unread expiry while progress arrives. Drain that
		// expiry before rearming, advance the epoch, and record the next arm's
		// scheduled deadline before Reset can make it observable. A stale expiry
		// consumed by the watchdog cannot then fence the reset attempt whether its
		// validity is captured before or after this reset.
		if !o.timer.Stop() {
			select {
			case <-o.timer.C():
			default:
			}
		}
		o.epoch++
		o.expiresAt = o.clock.Now().Add(o.deadline)
		o.timer.Reset(o.deadline)
	case runtime.DispositionTerminal:
		o.terminal = true
		o.timer.Stop()
	}
}

type attemptExpiry struct{ epoch uint64 }

// expiryForSignal binds a consumed timer signal to the arm that produced it.
// A progress reset advances expiresAt past a stale signal's fire time and bumps
// the epoch, so a signal received before that reset is rejected here even
// though its epoch was never captured against the old arm. The returned epoch
// continues to guard the post-capture race up to claimFence, where a progress
// reset that begins after this capture is also rejected.
func (o *attemptObserver) expiryForSignal(firedAt time.Time) (attemptExpiry, bool) {
	if o.beforeExpiryCapture != nil {
		o.beforeExpiryCapture()
	}
	o.mu.Lock()
	valid := !o.terminal && !o.fenced && !firedAt.Before(o.expiresAt)
	expiry := attemptExpiry{epoch: o.epoch}
	o.mu.Unlock()
	if o.afterExpiryCapture != nil {
		o.afterExpiryCapture(valid)
	}
	if !valid {
		return attemptExpiry{}, false
	}
	return expiry, true
}

func (o *attemptObserver) claimFence(expiry attemptExpiry) bool {
	o.mu.Lock()
	defer o.mu.Unlock()
	if o.terminal || o.fenced || o.epoch != expiry.epoch {
		return false
	}
	o.fenced = true
	o.timer.Stop()
	return true
}

func cloneLivenessMetadata(metadata map[string]string) map[string]string {
	cloned := make(map[string]string, len(metadata))
	for key, value := range metadata {
		cloned[key] = value
	}
	return cloned
}
func stalledRuntimeEvent(spec runtime.ExecutionSpec, obs stallObservation) runtime.RuntimeEvent {
	metadata := stallMetadata(spec.RunID, spec.Adapter, spec.Target, obs)
	return runtime.RuntimeEvent{RunID: spec.RunID, Type: runtime.EventTypeError, Timestamp: time.Now(), Error: "provider response stalled",
		Failure: &runtime.Failure{Code: runtime.FailureCodeResponseStalled, Message: "provider response stalled", Retryable: obs.fence == "confirmed", Metadata: cloneLivenessMetadata(metadata)}, Metadata: cloneLivenessMetadata(metadata)}
}
func stalledTunnelFrame(req runtime.ProviderTunnelRequest, obs stallObservation) runtime.ProviderTunnelFrame {
	metadata := stallMetadata(req.RunID, req.Adapter, req.Target, obs)
	return runtime.ProviderTunnelFrame{
		RunID:     req.RunID,
		TunnelID:  req.TunnelID,
		Kind:      runtime.ProviderTunnelFrameKindError,
		Error:     "provider response stalled",
		Timestamp: time.Now(),
		Failure: &runtime.Failure{
			Code:      runtime.FailureCodeResponseStalled,
			Message:   "provider response stalled",
			Retryable: obs.fence == "confirmed",
			Metadata:  cloneLivenessMetadata(metadata),
		},
		Metadata: cloneLivenessMetadata(metadata),
	}
}

// awaitAttempt owns the race between provider return, the request boundary,
// and the response-stall timer. A confirmed result means provider return was
// observed within close grace; otherwise resource cleanup remains provider-owned.
func awaitAttempt(
	execCtx context.Context,
	cancel context.CancelFunc,
	clock attemptClock,
	observer *attemptObserver,
	claimStall func(attemptExpiry) bool,
	providerDone <-chan error,
	probe healthProbe,
) attemptResult {
	for {
		select {
		case providerErr := <-providerDone:
			return attemptResult{providerErr: providerErr, providerReturned: true}
		case firedAt := <-observer.expired():
			expiry, valid := observer.expiryForSignal(firedAt)
			if !valid || !contextStillActive(execCtx) || !claimStall(expiry) {
				continue
			}
			cancel()
			return joinStallEvidence(clock, providerDone, probe)
		case <-execCtx.Done():
			cancel()
			return attemptResult{providerErr: <-providerDone, providerReturned: true}
		}
	}
}

func (n *Node) executeNormalizedAttempt(
	ctx, execCtx context.Context,
	cancel context.CancelFunc,
	adapter runtime.Provider,
	spec runtime.ExecutionSpec,
	ticket *admissionTicket,
	h *runHandle,
	sender protoSender,
	probe healthProbe,
	seq healthObservationSequencer,
) error {
	observer := newAttemptObserver(n.watchdogClock, time.Duration(spec.ResponseStallTimeoutMS)*time.Millisecond)
	sink := &terminalDeferringSink{
		inner:    &sessionSink{sess: sender, out: n.out, nodeID: n.nodeID, sessionID: normalizeSessionID(spec.SessionID), background: spec.Background},
		observer: observer,
	}
	providerDone := startProviderAttempt(func() error { return adapter.Execute(execCtx, spec, sink) })
	cleanup := newAttemptCleanup(func() {
		ticket.release()
		cancel()
		n.runs.deregister(spec.RunID)
		close(h.done)
	})
	result := awaitAttempt(execCtx, cancel, n.watchdogClock, observer, sink.claimStall, providerDone, probe)
	if !result.stalled {
		return n.finishNormalizedAttempt(ctx, spec, sink, cleanup, result.providerErr)
	}

	obs := stallObservationFrom(result, time.Duration(spec.ResponseStallTimeoutMS)*time.Millisecond, seq)
	n.liveness.Observe("normalized", obs)
	sink.queueClaimedTerminal(stalledRuntimeEvent(spec, obs))
	n.completeRun(spec, errProviderResponseStalled)
	if result.providerReturned {
		cleanup.run()
	} else {
		cleanup.afterProviderReturn(providerDone)
	}
	if err := sink.Flush(context.Background()); err != nil {
		n.logger.Warn("session: flush stalled terminal", zap.Error(err))
	}
	return errProviderResponseStalled
}

func (n *Node) finishNormalizedAttempt(
	ctx context.Context,
	spec runtime.ExecutionSpec,
	sink *terminalDeferringSink,
	cleanup *attemptCleanup,
	execErr error,
) error {
	cleanup.run()
	if !sink.hasTerminalObserved() {
		if synthErr := n.synthAndEmitTerminal(ctx, sink, spec, execErr); synthErr != nil && execErr == nil {
			execErr = synthErr
		}
	}
	n.completeRun(spec, execErr)
	if flushErr := sink.Flush(context.Background()); flushErr != nil && execErr == nil {
		return flushErr
	}
	return execErr
}

func (n *Node) consumeTunnelCredential(
	ctx context.Context,
	req *iop.ProviderTunnelRequest,
	tr *runtime.ProviderTunnelRequest,
) (*credentiallease.Material, error) {
	if n.credentialConsumer == nil && req.GetCredentialLease() == nil && req.GetCredentialBinding() == nil {
		return nil, nil
	}
	if n.credentialConsumer == nil || req.GetCredentialLease() == nil || req.GetCredentialBinding() == nil {
		return nil, errors.New("node: credential lease is required")
	}
	envelope, err := credentiallease.FromProto(req.GetCredentialLease())
	if err != nil {
		return nil, errors.New("node: credential lease rejected")
	}
	material, err := n.credentialConsumer.Consume(ctx, envelope, credentiallease.ExpectedFromProto(req.GetCredentialBinding()))
	if err != nil {
		return nil, errors.New("node: credential lease rejected")
	}
	tr.Credential = &runtime.ProviderCredential{HeaderName: material.HeaderName, Scheme: material.Scheme, Secret: material.Secret}
	return material, nil
}

func (n *Node) executeTunnelAttempt(
	execCtx context.Context,
	cancel context.CancelFunc,
	adapter runtime.ProviderTunnelAdapter,
	tr runtime.ProviderTunnelRequest,
	sink *tunnelSink,
	ticket *admissionTicket,
	h *runHandle,
	material *credentiallease.Material,
	probe healthProbe,
	seq healthObservationSequencer,
) error {
	providerDone := startProviderAttempt(func() error { return adapter.TunnelProvider(execCtx, tr, sink) })
	cleanup := newAttemptCleanup(func() {
		cancel()
		if tr.Credential != nil {
			tr.Credential.Zero()
		}
		if material != nil {
			material.Zero()
		}
		ticket.release()
		n.runs.deregister(tr.RunID)
		close(h.done)
	})
	result := awaitAttempt(execCtx, cancel, n.watchdogClock, sink.observer, sink.claimStall, providerDone, probe)
	if !result.stalled {
		cleanup.run()
		if result.providerErr != nil {
			n.logger.Warn("provider tunnel error", zap.String("run_id", tr.RunID), zap.String("tunnel_id", tr.TunnelID), zap.Error(result.providerErr))
		}
		return result.providerErr
	}

	if result.providerReturned {
		cleanup.run()
	} else {
		cleanup.afterProviderReturn(providerDone)
	}
	obs := stallObservationFrom(result, time.Duration(tr.ResponseStallTimeoutMS)*time.Millisecond, seq)
	n.liveness.Observe("provider_tunnel", obs)
	_ = sink.emitClaimedTerminal(context.Background(), stalledTunnelFrame(tr, obs))
	return errProviderResponseStalled
}

// terminalDeferringSink holds normalized terminal output until Node-local
// admission has released its slot. emitMu is the single emission authority for
// accepted provider events and a watchdog fence claim.
type terminalDeferringSink struct {
	inner    runtime.EventSink
	observer *attemptObserver

	emitMu           sync.Mutex
	mu               sync.Mutex
	deferring        bool
	terminalObserved bool
	fenced           bool
	deferred         []runtime.RuntimeEvent

	// beforeStallClaim is a deterministic ordering seam for package tests.
	beforeStallClaim func()
	afterStallClaim  func(bool)
}

func (s *terminalDeferringSink) Emit(ctx context.Context, event runtime.RuntimeEvent) error {
	s.emitMu.Lock()
	defer s.emitMu.Unlock()
	s.mu.Lock()
	if s.terminalObserved || s.fenced {
		s.mu.Unlock()
		return nil
	}
	if s.observer != nil {
		s.observer.observe(runtime.ClassifyRuntimeEvent(event))
	}
	if runtime.IsTerminalEvent(event.Type) {
		s.terminalObserved = true
	}
	if s.deferring || runtime.IsTerminalEvent(event.Type) {
		s.deferring = true
		s.deferred = append(s.deferred, event)
		s.mu.Unlock()
		return nil
	}
	s.mu.Unlock()
	return s.inner.Emit(ctx, event)
}

func (s *terminalDeferringSink) claimStall(expiry attemptExpiry) bool {
	if s.beforeStallClaim != nil {
		s.beforeStallClaim()
	}
	s.emitMu.Lock()
	s.mu.Lock()
	if s.terminalObserved || s.fenced || (s.observer != nil && !s.observer.claimFence(expiry)) {
		s.mu.Unlock()
		s.emitMu.Unlock()
		if s.afterStallClaim != nil {
			s.afterStallClaim(false)
		}
		return false
	}
	s.fenced, s.terminalObserved, s.deferring = true, true, true
	s.mu.Unlock()
	s.emitMu.Unlock()
	if s.afterStallClaim != nil {
		s.afterStallClaim(true)
	}
	return true
}

func (s *terminalDeferringSink) queueClaimedTerminal(event runtime.RuntimeEvent) {
	s.emitMu.Lock()
	defer s.emitMu.Unlock()
	s.mu.Lock()
	s.deferred = append(s.deferred, event)
	s.mu.Unlock()
}

func (s *terminalDeferringSink) Flush(ctx context.Context) error {
	s.emitMu.Lock()
	defer s.emitMu.Unlock()
	s.mu.Lock()
	events := append([]runtime.RuntimeEvent(nil), s.deferred...)
	s.deferred = nil
	s.deferring = false
	s.mu.Unlock()
	for _, event := range events {
		if err := s.inner.Emit(ctx, event); err != nil {
			return err
		}
	}
	return nil
}

func (s *terminalDeferringSink) hasTerminalObserved() bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.terminalObserved
}

// tunnelSink holds its emission lock through Send. A watchdog fence therefore
// cannot overtake a frame that was accepted before the fence claim.
type tunnelSink struct {
	sess      protoSender
	nodeID    string
	nodeAlias string
	observer  *attemptObserver
	mu        sync.Mutex
	terminal  bool
	fenced    bool

	// beforeStallClaim is a deterministic ordering seam for package tests.
	beforeStallClaim func()
	afterStallClaim  func(bool)
}

func (s *tunnelSink) EmitTunnelFrame(ctx context.Context, frame runtime.ProviderTunnelFrame) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.terminal || s.fenced {
		return nil
	}
	disposition := runtime.ClassifyProviderTunnelFrame(frame)
	if s.observer != nil {
		s.observer.observe(disposition)
	}
	if disposition == runtime.DispositionTerminal {
		s.terminal = true
	}
	return s.emit(ctx, frame)
}

func (s *tunnelSink) claimStall(expiry attemptExpiry) bool {
	if s.beforeStallClaim != nil {
		s.beforeStallClaim()
	}
	s.mu.Lock()
	if s.terminal || s.fenced || (s.observer != nil && !s.observer.claimFence(expiry)) {
		s.mu.Unlock()
		if s.afterStallClaim != nil {
			s.afterStallClaim(false)
		}
		return false
	}
	s.fenced, s.terminal = true, true
	s.mu.Unlock()
	if s.afterStallClaim != nil {
		s.afterStallClaim(true)
	}
	return true
}

func (s *tunnelSink) emitClaimedTerminal(ctx context.Context, frame runtime.ProviderTunnelFrame) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.emit(ctx, frame)
}

func (s *tunnelSink) emit(ctx context.Context, frame runtime.ProviderTunnelFrame) error {
	tf := tunnelFrameToProto(frame, s.nodeID, s.nodeAlias)
	if s.sess != nil {
		return s.sess.Send(tf)
	}
	return nil
}

func tunnelFrameToProto(frame runtime.ProviderTunnelFrame, nodeID, nodeAlias string) *iop.ProviderTunnelFrame {
	var usage *iop.Usage
	if frame.Usage != nil {
		usage = &iop.Usage{InputTokens: int32(frame.Usage.InputTokens), OutputTokens: int32(frame.Usage.OutputTokens), ReasoningTokens: int32(frame.Usage.ReasoningTokens), CachedInputTokens: int32(frame.Usage.CachedInputTokens)}
	}
	protoKind := iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_UNSPECIFIED
	switch frame.Kind {
	case runtime.ProviderTunnelFrameKindResponseStart:
		protoKind = iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START
	case runtime.ProviderTunnelFrameKindBody:
		protoKind = iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY
	case runtime.ProviderTunnelFrameKindEnd:
		protoKind = iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END
	case runtime.ProviderTunnelFrameKindError:
		protoKind = iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR
	case runtime.ProviderTunnelFrameKindUsage:
		protoKind = iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_USAGE
	}
	return &iop.ProviderTunnelFrame{
		RunId: frame.RunID, TunnelId: frame.TunnelID, Sequence: frame.Sequence, Kind: protoKind,
		StatusCode: int32(frame.StatusCode), Headers: frame.Headers, Body: frame.Body, End: frame.End,
		Error: frame.Error, Failure: executionFailureToProto(frame.Failure), Usage: usage, Metadata: cloneLivenessMetadata(frame.Metadata), Timestamp: frame.Timestamp.UnixNano(),
		NodeId: nodeID, NodeAlias: nodeAlias,
	}
}
