package agent

import (
	"context"
	"errors"
	"fmt"
	"io"
	"sync"
	"time"

	rvboxv1 "github.com/rvbox/rvbox/gen/go/rvbox/v1"
	"github.com/rvbox/rvbox/internal/client/spool"
	"github.com/rvbox/rvbox/internal/client/supervisor"
	"github.com/rvbox/rvbox/internal/domain"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/timestamppb"
)

// Executor bridges durable dispatch records to the platform supervisor. It
// owns no network state: a command's process, pipes, and spool writes continue
// after the WebSocket session is replaced.
type Executor struct {
	Store      *spool.Store
	Supervisor supervisor.Supervisor
	WorkDir    string
	Now        func() time.Time
	Notify     func(domain.UUID)

	mu     sync.Mutex
	active map[domain.UUID]supervisor.Process
	cancel map[domain.UUID]context.CancelFunc
	forced map[domain.UUID]bool
}

type ExecutorOptions struct {
	Store      *spool.Store
	Supervisor supervisor.Supervisor
	WorkDir    string
	Now        func() time.Time
	Notify     func(domain.UUID)
}

func NewExecutor(options ExecutorOptions) (*Executor, error) {
	if options.Store == nil || options.Supervisor == nil || options.WorkDir == "" {
		return nil, errors.New("executor requires durable store, supervisor, and work directory")
	}
	if options.Now == nil {
		options.Now = func() time.Time { return time.Now().UTC() }
	}
	return &Executor{Store: options.Store, Supervisor: options.Supervisor, WorkDir: options.WorkDir, Now: options.Now, Notify: options.Notify, active: make(map[domain.UUID]supervisor.Process), cancel: make(map[domain.UUID]context.CancelFunc), forced: make(map[domain.UUID]bool)}, nil
}

// Dispatch is safe to invoke after CommandAccepted has been sent. Script
// commands intentionally wait for ScriptCommit; the server may deliver the
// descriptor and body in separate frames.
func (executor *Executor) Dispatch(ctx context.Context, _ Session, dispatch *rvboxv1.CommandDispatch) error {
	if executor == nil || dispatch == nil {
		return errors.New("invalid command dispatch")
	}
	issue, err := domain.ParseUUIDv7(dispatch.GetIssueUuid())
	if err != nil {
		return err
	}
	if dispatch.GetSpec().GetScript() != nil {
		if _, err := executor.Store.ScriptBody(ctx, issue); errors.Is(err, spool.ErrScriptNotReady) {
			return nil
		} else if err != nil {
			return executor.reject(ctx, issue, dispatch.GetCommandRevision(), err)
		}
	}
	return executor.launch(ctx, issue, dispatch.GetCommandRevision(), dispatch.GetSpec())
}

// ScriptReady launches a script after the durable commit barrier. Repeated
// progress/commit frames are harmless because the active map and spool phase
// make launch at-most-once.
func (executor *Executor) ScriptReady(ctx context.Context, _ Session, issue domain.UUID) error {
	command, err := executor.Store.GetCommand(ctx, issue)
	if err != nil {
		return err
	}
	var spec rvboxv1.ExecutionSpec
	if err := proto.Unmarshal(command.ExecutionSpec, &spec); err != nil {
		return executor.reject(ctx, issue, command.Revision, err)
	}
	if spec.GetScript() == nil {
		return nil
	}
	return executor.launch(ctx, issue, command.Revision, &spec)
}

func (executor *Executor) launch(ctx context.Context, issue domain.UUID, revision uint64, spec *rvboxv1.ExecutionSpec) error {
	if spec == nil || revision == 0 {
		return executor.reject(ctx, issue, revision, errors.New("missing execution specification"))
	}
	executor.mu.Lock()
	if _, exists := executor.active[issue]; exists {
		executor.mu.Unlock()
		return nil
	}
	executor.mu.Unlock()
	var scriptBody []byte
	if spec.GetScript() != nil {
		body, err := executor.Store.ScriptBody(ctx, issue)
		if err != nil {
			if errors.Is(err, spool.ErrScriptNotReady) {
				return nil
			}
			return executor.reject(ctx, issue, revision, err)
		}
		scriptBody = body
	}
	workingDirectory := spec.GetCwd()
	if workingDirectory == "" {
		workingDirectory = executor.WorkDir
	}
	if err := executor.Store.SetLaunchPhase(ctx, issue, domain.LaunchPhasePrepared, "", 0); err != nil {
		return err
	}
	// The native Windows supervisor creates/assigns the Job while the child is
	// suspended and releases it before returning. Marking authorization before
	// that call makes a daemon crash in any of those windows recover as an
	// interrupted, non-redispatchable command.
	if err := executor.Store.SetLaunchPhase(ctx, issue, domain.LaunchPhaseAuthorized, "pending", 0); err != nil {
		return err
	}
	runContext, cancel := context.WithCancel(context.Background())
	process, err := executor.Supervisor.Start(runContext, supervisor.StartSpec{IssueUUID: issue, CommandRevision: revision, Execution: proto.Clone(spec).(*rvboxv1.ExecutionSpec), ScriptBody: scriptBody, WorkingDirectory: workingDirectory, Environment: cloneEnvironment(spec.GetEnvOverrides()), ExecutionProfiles: executionProfileNames(spec.GetExecutionProfiles())})
	if err != nil {
		cancel()
		_ = executor.Store.SetLaunchPhase(context.Background(), issue, domain.LaunchPhaseNone, "", 0)
		var startErr *supervisor.StartError
		if errors.As(err, &startErr) {
			return executor.rejectWithIdentity(ctx, issue, revision, err, startErr.WindowsIdentity)
		}
		return executor.reject(ctx, issue, revision, err)
	}
	identity := process.Identity()
	if err := executor.Store.SetLaunchPhase(ctx, issue, domain.LaunchPhaseAuthorized, identity.Context, 0); err != nil {
		_, _ = executor.Supervisor.Signal(context.Background(), process, supervisor.SignalKill)
		cancel()
		return err
	}
	executor.mu.Lock()
	executor.active[issue] = process
	executor.cancel[issue] = cancel
	executor.mu.Unlock()
	if _, err := executor.Store.AppendLifecycleWithIdentity(ctx, issue, uint32(rvboxv1.CommandLifecycle_COMMAND_RUNNING), revision, "process started", executor.Now(), identity.WindowsIdentity()); err != nil {
		_, _ = executor.Supervisor.Signal(context.Background(), process, supervisor.SignalKill)
		cancel()
		executor.remove(issue)
		return err
	}
	executor.notify(issue)
	go executor.watch(runContext, issue, revision, process, cancel)
	return nil
}

func (executor *Executor) watch(ctx context.Context, issue domain.UUID, revision uint64, process supervisor.Process, cancel context.CancelFunc) {
	defer cancel()
	for {
		chunk, err := process.ReadOutput(ctx)
		if errors.Is(err, io.EOF) {
			break
		}
		if err != nil {
			_ = executor.appendIncomplete(context.Background(), issue, revision, err.Error())
			break
		}
		if len(chunk.Data) == 0 {
			continue
		}
		if _, err := executor.Store.AppendOutput(context.Background(), issue, spool.OutputInput{Stream: chunk.Stream, Raw: chunk.Data, ObservedAt: executor.Now()}); err != nil {
			_ = executor.appendIncomplete(context.Background(), issue, revision, err.Error())
			_, _ = executor.Supervisor.Signal(context.Background(), process, supervisor.SignalKill)
			break
		}
		executor.notify(issue)
	}
	status, waitErr := process.Wait(context.Background())
	phase := rvboxv1.CommandLifecycle_COMMAND_FAILED
	if waitErr == nil && status.Code == 0 {
		phase = rvboxv1.CommandLifecycle_COMMAND_SUCCEEDED
	}
	executor.mu.Lock()
	if executor.forced[issue] {
		phase = rvboxv1.CommandLifecycle_COMMAND_TERMINATED
	}
	executor.mu.Unlock()
	detail := fmt.Sprintf("exit code %d", status.Code)
	if waitErr != nil {
		detail = boundedError(waitErr)
	}
	if err := executor.appendLifecycleWithIdentity(context.Background(), issue, revision, phase, detail, process.Identity().WindowsIdentity()); err == nil {
		executor.notify(issue)
	}
	executor.remove(issue)
}

func (executor *Executor) appendLifecycle(ctx context.Context, issue domain.UUID, revision uint64, phase rvboxv1.CommandLifecycle, detail string) error {
	return executor.appendLifecycleWithIdentity(ctx, issue, revision, phase, detail, nil)
}

func (executor *Executor) appendLifecycleWithIdentity(ctx context.Context, issue domain.UUID, revision uint64, phase rvboxv1.CommandLifecycle, detail string, identity *rvboxv1.WindowsExecutionIdentity) error {
	_, err := executor.Store.AppendLifecycleWithIdentity(ctx, issue, uint32(phase), revision, detail, executor.Now(), identity)
	return err
}

func (executor *Executor) reject(ctx context.Context, issue domain.UUID, revision uint64, cause error) error {
	return executor.rejectWithIdentity(ctx, issue, revision, cause, nil)
}

func (executor *Executor) rejectWithIdentity(ctx context.Context, issue domain.UUID, revision uint64, cause error, identity *rvboxv1.WindowsExecutionIdentity) error {
	if revision == 0 {
		return cause
	}
	if err := executor.appendLifecycleWithIdentity(ctx, issue, revision, rvboxv1.CommandLifecycle_COMMAND_REJECTED, boundedError(cause), identity); err != nil {
		return err
	}
	executor.notify(issue)
	return nil
}

func (executor *Executor) appendIncomplete(ctx context.Context, issue domain.UUID, _ uint64, detail string) error {
	payload, err := proto.MarshalOptions{Deterministic: true}.Marshal(&rvboxv1.CommandEvent{IssueUuid: issue.String(), ObservedAt: timestamppb.New(executor.Now()), Payload: &rvboxv1.CommandEvent_OutputIncomplete{OutputIncomplete: &rvboxv1.OutputIncomplete{Reason: boundedError(errors.New(detail))}}})
	if err != nil {
		return err
	}
	_, err = executor.Store.AppendEvent(ctx, issue, spool.EventInput{Kind: 11, Compression: 1, RawBytes: uint64(len(payload)), Payload: payload, CreatedAt: executor.Now()})
	if err == nil {
		executor.notify(issue)
	}
	return err
}

func (executor *Executor) Stdin(ctx context.Context, _ Session, input *rvboxv1.StdinWrite) error {
	if input == nil {
		return errors.New("missing stdin request")
	}
	issue, err := domain.ParseUUIDv7(input.GetIssueUuid())
	if err != nil {
		return err
	}
	executor.mu.Lock()
	process := executor.active[issue]
	executor.mu.Unlock()
	if process == nil {
		return executor.appendStdinAck(ctx, issue, input.GetWriteSeq(), false, "command is not running")
	}
	err = process.WriteStdin(ctx, input.GetData(), input.GetAppendNewline())
	detail := ""
	if err != nil {
		detail = boundedError(err)
	}
	return executor.appendStdinAck(ctx, issue, input.GetWriteSeq(), err == nil, detail)
}

func (executor *Executor) CloseStdin(ctx context.Context, _ Session, input *rvboxv1.CloseStdin) error {
	if input == nil {
		return errors.New("missing close-stdin request")
	}
	issue, err := domain.ParseUUIDv7(input.GetIssueUuid())
	if err != nil {
		return err
	}
	executor.mu.Lock()
	process := executor.active[issue]
	executor.mu.Unlock()
	if process == nil {
		return executor.appendStdinAck(ctx, issue, input.GetWriteSeq(), false, "command is not running")
	}
	err = process.CloseStdin(ctx)
	detail := ""
	if err != nil {
		detail = boundedError(err)
	}
	return executor.appendStdinAck(ctx, issue, input.GetWriteSeq(), err == nil, detail)
}

func (executor *Executor) appendStdinAck(ctx context.Context, issue domain.UUID, writeSeq uint64, accepted bool, detail string) error {
	payload, err := proto.MarshalOptions{Deterministic: true}.Marshal(&rvboxv1.CommandEvent{IssueUuid: issue.String(), ObservedAt: timestamppb.New(executor.Now()), Payload: &rvboxv1.CommandEvent_StdinAck{StdinAck: &rvboxv1.StdinAcknowledgement{WriteSeq: writeSeq, Detail: detail, StdinClosed: !accepted}}})
	if err != nil {
		return err
	}
	_, err = executor.Store.AppendEvent(ctx, issue, spool.EventInput{Kind: 7, Compression: 1, RawBytes: uint64(len(payload)), Payload: payload, CreatedAt: executor.Now(), UseCloseout: false})
	if err == nil {
		executor.notify(issue)
	}
	return err
}

func (executor *Executor) Signal(ctx context.Context, _ Session, input *rvboxv1.SignalCommand) error {
	if input == nil {
		return errors.New("missing signal request")
	}
	issue, err := domain.ParseUUIDv7(input.GetIssueUuid())
	if err != nil {
		return err
	}
	executor.mu.Lock()
	process := executor.active[issue]
	executor.mu.Unlock()
	if process == nil {
		return errors.New("command is not running")
	}
	kind := supervisor.SignalTerm
	if input.GetSignal() == rvboxv1.SignalKind_SIGNAL_KILL {
		kind = supervisor.SignalKill
	} else if input.GetSignal() != rvboxv1.SignalKind_SIGNAL_TERM {
		return errors.New("unsupported Windows signal")
	}
	outcome, err := executor.Supervisor.Signal(ctx, process, kind)
	if kind == supervisor.SignalKill || outcome.Escalated {
		executor.mu.Lock()
		executor.forced[issue] = true
		executor.mu.Unlock()
	}
	return executor.appendSignalResult(ctx, issue, input.GetCommandRevision(), input.GetSignal(), err == nil && outcome.Delivered, outcome, err)
}

func (executor *Executor) appendSignalResult(ctx context.Context, issue domain.UUID, revision uint64, signal rvboxv1.SignalKind, accepted bool, outcome supervisor.SignalOutcome, cause error) error {
	detail := outcome.Detail
	if cause != nil {
		detail = boundedError(cause)
	}
	payload, err := proto.MarshalOptions{Deterministic: true}.Marshal(&rvboxv1.CommandEvent{IssueUuid: issue.String(), ObservedAt: timestamppb.New(executor.Now()), Payload: &rvboxv1.CommandEvent_SignalResult{SignalResult: &rvboxv1.SignalResult{Signal: signal, Accepted: accepted, GracefulDeliveryAttempted: signal == rvboxv1.SignalKind_SIGNAL_TERM, ForcedTerminationUsed: outcome.Escalated || signal == rvboxv1.SignalKind_SIGNAL_KILL, Detail: detail, CommandRevision: revision}}})
	if err != nil {
		return err
	}
	_, err = executor.Store.AppendEvent(ctx, issue, spool.EventInput{Kind: 8, Compression: 1, RawBytes: uint64(len(payload)), Payload: payload, CreatedAt: executor.Now()})
	if err == nil {
		executor.notify(issue)
	}
	return err
}

func (executor *Executor) Terminate(ctx context.Context, issue domain.UUID) error {
	executor.mu.Lock()
	process := executor.active[issue]
	executor.forced[issue] = true
	executor.mu.Unlock()
	if process == nil {
		return nil
	}
	returnError := error(nil)
	if _, err := executor.Supervisor.Signal(ctx, process, supervisor.SignalKill); err != nil {
		returnError = err
	}
	return returnError
}

func (executor *Executor) StopAll(ctx context.Context) error {
	return executor.Supervisor.StopAll(ctx)
}

func (executor *Executor) remove(issue domain.UUID) {
	executor.mu.Lock()
	delete(executor.active, issue)
	delete(executor.cancel, issue)
	delete(executor.forced, issue)
	executor.mu.Unlock()
}

func (executor *Executor) notify(issue domain.UUID) {
	if executor.Notify != nil {
		executor.Notify(issue)
	}
}

func cloneEnvironment(input map[string]string) map[string]string {
	if input == nil {
		return nil
	}
	result := make(map[string]string, len(input))
	for key, value := range input {
		result[key] = value
	}
	return result
}

func executionProfileNames(input []rvboxv1.ExecutionProfile) []string {
	result := make([]string, 0, len(input))
	for _, profile := range input {
		result = append(result, profile.String())
	}
	return result
}