package agent

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/drivercontract"
)

// EncodeStdinMessage formats a routed message as a pi RPC `prompt` command.
// Pi's primary delivery path is DeliverRuntimeInput (protocol_direct); this
// fallback is used only by callers without a structured runtime-input driver.
func (d *PiRPCDriver) EncodeStdinMessage(text, sessionID string) []byte {
	cmd := map[string]any{"type": "prompt", "message": text, "streamingBehavior": "followUp"}
	data, _ := json.Marshal(cmd)
	return append(data, '\n')
}

// DeliverBusyStdinInput accepts compiler-produced input while a root run is
// active via pi's steer queue.
func (d *PiRPCDriver) DeliverBusyStdinInput(ctx context.Context, input drivercontract.RuntimeInput, sessionID string) error {
	if ctx == nil {
		ctx = context.Background()
	}
	return d.DeliverRuntimeInput(ctx, input, sessionID)
}

// DeliverRuntimeInput sends a user message to pi as an RPC `prompt` command and
// returns on command acceptance (pi's preflight response). The root run stays
// active until agent_settled settles the same generation; the caller must not
// treat a successful return as turn completion.
func (d *PiRPCDriver) DeliverRuntimeInput(ctx context.Context, input drivercontract.RuntimeInput, sessionID string) error {
	if input.Empty() {
		return nil
	}
	if ctx == nil {
		ctx = context.Background()
	}
	if err := d.validateSession(sessionID); err != nil {
		return err
	}
	if err := ctx.Err(); err != nil {
		// The caller canceled before the prompt could be registered or
		// serialized. This is a definitive pre-mutation rejection, so a
		// same-generation reconcile may safely return the input to queued.
		d.rememberDefinitiveInputRejection(input.ID, sessionID)
		return fmt.Errorf("%w: %w", errRuntimeInputMutationNotAttempted, err)
	}
	userInput, err := piRPCUserInput(input)
	if err != nil {
		return err
	}
	// A prior authoritative rejection/idle snapshot may have returned this same
	// durable id to the queue. A new mutation generation must not inherit that
	// negative fence; any definite failure below installs it again.
	d.forgetDefinitiveInputRejection(input.ID, sessionID)
	d.mu.Lock()
	if d.rootRunActive {
		d.mu.Unlock()
		return fmt.Errorf("pi: a root run is already active")
	}
	if d.promptAccepted != nil {
		d.mu.Unlock()
		return fmt.Errorf("pi: prompt already pending")
	}
	accepted := make(chan error, 1)
	d.promptAccepted = accepted
	d.promptInputID = strings.TrimSpace(input.ID)
	runtimeCtx := d.runtimeCtx
	if runtimeCtx == nil {
		runtimeCtx = context.Background()
	}
	d.mu.Unlock()

	command := map[string]any{
		"type":              "prompt",
		"message":           userInput,
		"streamingBehavior": "followUp",
	}
	go func() {
		if err := d.sendCommand(ctx, command); err != nil {
			if !piRPCMutationAmbiguous(err) {
				d.rememberDefinitiveInputRejection(input.ID, sessionID)
			}
			d.resolvePromptAcceptance(err)
			return
		}
		// The prompt response handler (handleLine) starts the generation-fenced
		// root run synchronously when the response is routed, so event frames
		// that follow on stdout observe rootRunActive=true. This goroutine only
		// resolves the acceptance handshake.
		d.resolvePromptAcceptance(nil)
	}()
	select {
	case <-ctx.Done():
		return errPiRPCTransportClosed
	case err := <-accepted:
		return err
	}
}

// waitForInitialAck closes the startup gap between Pi's prompt RPC
// response and the durable/user-input boundary represented by the correlated
// marker. A successful prompt response alone must not let Process publish the
// first message as acknowledged.
func (d *PiRPCDriver) waitForInitialAck(ctx context.Context, inputID string) error {
	inputID = strings.TrimSpace(inputID)
	if inputID == "" {
		return nil
	}
	if ctx == nil {
		ctx = context.Background()
	}
	d.mu.Lock()
	if d.ackedInputIDs[inputID] {
		d.mu.Unlock()
		return nil
	}
	if d.transportClosed {
		d.mu.Unlock()
		return errPiRPCTransportClosed
	}
	waiter := make(chan error, 1)
	if d.initialAckWaiters == nil {
		d.initialAckWaiters = make(map[string]chan error)
	}
	d.initialAckWaiters[inputID] = waiter
	settled := d.rootRunSettled
	d.mu.Unlock()

	select {
	case err := <-waiter:
		return err
	case <-settled:
		d.mu.Lock()
		acked := d.ackedInputIDs[inputID]
		rootError := strings.TrimSpace(d.rootError)
		if current, ok := d.initialAckWaiters[inputID]; ok && current == waiter {
			delete(d.initialAckWaiters, inputID)
		}
		d.mu.Unlock()
		if acked {
			return nil
		}
		if rootError != "" {
			return fmt.Errorf("pi: initial prompt settled before acknowledgement: %s", rootError)
		}
		return fmt.Errorf("pi: initial prompt settled before runtime input acknowledgement")
	case <-ctx.Done():
		d.mu.Lock()
		if current, ok := d.initialAckWaiters[inputID]; ok && current == waiter {
			delete(d.initialAckWaiters, inputID)
		}
		d.mu.Unlock()
		return ctx.Err()
	}
}

// beginRootRun starts a new generation-fenced root run. Streaming/dedup state
// is reset so late events from a prior generation cannot corrupt it.
func (d *PiRPCDriver) beginRootRun() {
	d.mu.Lock()
	d.rootRunGeneration++
	d.rootRunTurnID = fmt.Sprintf("pi-gen-%d", d.rootRunGeneration)
	d.rootRunActive = true
	d.rootRunSettled = make(chan struct{})
	d.rootStartEmitted = false
	d.rootStopReason = ""
	d.rootAborted = false
	d.rootError = ""
	d.rootErrorEmitted = false
	d.agentRunActive = false
	d.agentRunCount = 0
	d.turnActive = false
	d.turnCount = 0
	d.assistant = nil
	d.ackedInputIDs = make(map[string]bool)
	d.reconciledAckedInputIDs = make(map[string]bool)
	d.finalizedAssistantIDs = make(map[string]bool)
	d.mu.Unlock()
}

// CurrentActiveTurnID returns the current active root-run turn identity, or ""
// when no root run is active. Implemented for drivercontract.ActiveTurnIDProvider
// so Process.runtimeBoundarySuperseded rejects stale old-generation terminal
// events once a replacement generation is active.
func (d *PiRPCDriver) CurrentActiveTurnID() string {
	d.mu.Lock()
	defer d.mu.Unlock()
	if !d.rootRunActive {
		return ""
	}
	return d.rootRunTurnID
}

// resolvePromptAcceptance completes the DeliverRuntimeInput acceptance
// handshake.
func (d *PiRPCDriver) resolvePromptAcceptance(err error) {
	d.mu.Lock()
	accepted := d.promptAccepted
	d.promptAccepted = nil
	// Keep the correlated id until the user message marker is observed. The
	// prompt response is only an acceptance boundary; clearing this here would
	// make a marker that arrives after the response (or after agent_settled)
	// indistinguishable from an unsolicited event and would let reconciliation
	// incorrectly call a still-owned input absent.
	if err != nil {
		d.promptInputID = ""
	}
	d.mu.Unlock()
	if accepted != nil {
		accepted <- err
	}
}

// DeliverAckedBusyRuntimeInput queues a correlated steer message while a root
// run is active and returns as soon as the steer mutation transport is ACCEPTED
// by the runtime (the `steer` RPC response). The marked user message is
// acknowledged later by an independent asynchronous runtime_input_acked event; the
// caller's persisted dispatching owner stays pending-only until then.
//
// Correlation is registered before send and retained until ackInputIfMarked,
// transport retirement, or bounded lifecycle cleanup — never removed on the
// success path. A definite RPC rejection removes correlation and returns the
// error so Process can requeue; an ambiguous transport outcome preserves
// ownership (fail-close) and never retries an accepted steer.
func (d *PiRPCDriver) DeliverAckedBusyRuntimeInput(ctx context.Context, input drivercontract.RuntimeInput, sessionID string) error {
	if input.Empty() {
		return nil
	}
	if strings.TrimSpace(input.ID) == "" {
		return fmt.Errorf("pi: correlated steer input id is required")
	}
	if err := d.validateSession(sessionID); err != nil {
		return err
	}
	userInput, err := piRPCUserInput(input)
	if err != nil {
		return err
	}
	inputID := strings.TrimSpace(input.ID)

	// Register correlation before the send so a marked user message that races
	// ahead of the steer response is still recognized and finalized exactly once.
	d.mu.Lock()
	if d.steerInFlight == nil {
		d.steerInFlight = make(map[string]bool)
	}
	if d.steerInFlight[inputID] {
		d.mu.Unlock()
		// The durable Process owner may be rebuilt/requeued while this live Pi
		// driver still owns an earlier accepted (or transport-ambiguous) steer.
		// Treat the repeated admission as idempotent acceptance: sending the
		// mutation again could duplicate the user message, while returning an
		// ordinary error makes the daemon's transient retry agentteam exhaust on its
		// own retained correlation. Reconciliation remains responsible for
		// deciding whether the one mutation was acknowledged or authoritatively absent.
		return nil
	}
	if d.startupSteerIdleObservedAt == nil {
		d.startupSteerIdleObservedAt = make(map[string]time.Time)
	}
	delete(d.startupSteerIdleObservedAt, inputID)
	d.steerInFlight[inputID] = true
	d.mu.Unlock()

	command := map[string]any{"type": "steer", "message": userInput}
	if err := d.sendCommand(ctx, command); err != nil {
		// If the marked message was already acknowledged, the runtime accepted the steer
		// despite the RPC error: treat as accepted.
		d.mu.Lock()
		acked := !d.steerInFlight[inputID]
		d.mu.Unlock()
		if acked {
			return nil
		}
		// Ambiguous transport outcome: keep ownership fail-closed (never retry
		// an accepted steer), remove nothing; Process classifies the error.
		if piRPCMutationAmbiguous(err) {
			return errors.Join(errPiRPCTransportAmbiguous, err)
		}
		// Definite rejection: remove correlation and return the error so
		// Process can requeue the delivery.
		d.mu.Lock()
		delete(d.steerInFlight, inputID)
		d.mu.Unlock()
		d.rememberDefinitiveInputRejection(inputID, sessionID)
		return fmt.Errorf("pi: steer: %w", err)
	}
	// Successful steer response: return acceptance immediately. Correlation stays
	// registered for the later runtime_input_acked event.
	return nil
}

// InterruptActiveRun aborts the current root run and waits for the same
// generation's authoritative agent_settled boundary, not an earlier Pi
// turn/agent event.
func (d *PiRPCDriver) InterruptActiveRun(ctx context.Context, sessionID string) error {
	if err := d.validateSession(sessionID); err != nil {
		return err
	}
	d.mu.Lock()
	settled := d.rootRunSettled
	active := d.rootRunActive
	d.mu.Unlock()
	if !active {
		return nil
	}
	// Mark the generation aborted BEFORE dispatching so the readLoop cannot race
	// an agent_settled past the flag; the final settlement is aborted unless a
	// real assistant stop reason supersedes it.
	d.mu.Lock()
	d.rootAborted = true
	d.mu.Unlock()
	if err := d.sendCommand(ctx, map[string]any{"type": "abort"}); err != nil {
		select {
		case <-settled:
			return nil
		default:
		}
		return fmt.Errorf("pi: abort: %w", err)
	}
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-settled:
		return nil
	}
}

// DeliverInterruptingRuntimeInput aborts the active root run (if any), waits
// for the same generation's agent_settled, then delivers the replacement input
// as a new root run.
func (d *PiRPCDriver) DeliverInterruptingRuntimeInput(ctx context.Context, input drivercontract.RuntimeInput, sessionID string) error {
	d.mu.Lock()
	settled := d.rootRunSettled
	active := d.rootRunActive
	d.mu.Unlock()
	if active {
		if err := d.InterruptActiveRun(ctx, sessionID); err != nil {
			return err
		}
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-settled:
		}
	}
	return d.DeliverRuntimeInput(ctx, input, sessionID)
}

func (d *PiRPCDriver) validateSession(sessionID string) error {
	d.mu.Lock()
	current := d.sessionID
	d.mu.Unlock()
	sessionID = strings.TrimSpace(sessionID)
	if sessionID == "" {
		sessionID = current
	}
	if current == "" {
		return fmt.Errorf("pi: session is not initialized")
	}
	if sessionID != current {
		return fmt.Errorf("pi: session mismatch: got %q, want %q", sessionID, current)
	}
	return nil
}