package runtime

import (
	"context"
	"fmt"
	"runtime"
	"strings"
	"sync"
	"time"

	"browser.local/run/config"
	"browser.local/run/protocol"
	"browser.local/run/spool"
)

type WorkerClient interface {
	Hello(context.Context, protocol.RunHelloRequest) (protocol.RunHelloResponse, error)
	Heartbeat(context.Context, protocol.RunHeartbeatRequest) (protocol.RunHeartbeatResponse, error)
	ClaimJob(context.Context, protocol.RunJobClaimRequest) (protocol.RunJobClaimResponse, error)
	AckJob(context.Context, protocol.RunJobAckRequest) (protocol.RunJobAckResponse, error)
	UpdateJobProgress(context.Context, protocol.RunJobProgressRequest) (protocol.RunJobProgressResponse, error)
	CompleteJob(context.Context, protocol.RunJobResultRequest) (protocol.RunJobResultResponse, error)
	PollJobCancel(context.Context, protocol.RunJobCancelPollRequest) (protocol.RunJobCancelPollResponse, error)
	ReconcileJobs(context.Context, protocol.RunJobReconcileRequest) (protocol.RunJobReconcileResponse, error)
}

type Worker struct {
	cfg      config.Config
	client   WorkerClient
	executor LifecycleExecutor
	state    WorkerState
	journal  *JobJournal
}

type WorkerState struct {
	RunEndpointID string
	SessionToken  string
	Capabilities  []string
	Capacity      protocol.RunCapacityReport
	LastHeartbeat time.Time
	Sequence      uint64
}

func NewWorker(cfg config.Config, client WorkerClient, options ...LifecycleExecutorOption) (*Worker, error) {
	if client == nil {
		return nil, fmt.Errorf("worker client is required")
	}
	if cfg.RunEndpointID == "" {
		cfg.RunEndpointID = config.DefaultEndpointID
	}
	if cfg.DisplayName == "" {
		cfg.DisplayName = config.DefaultDisplayName
	}
	if cfg.Version == "" {
		cfg.Version = config.DefaultVersion
	}
	if cfg.MaxJobs <= 0 {
		cfg.MaxJobs = 1
	}
	executorOptions := append([]LifecycleExecutorOption{
		WithLifecycleWorkspaceRoot(cfg.WorkspaceRoot),
	}, options...)
	return &Worker{
		cfg:      cfg,
		client:   client,
		executor: NewLifecycleExecutor(executorOptions...),
		state: WorkerState{
			RunEndpointID: cfg.RunEndpointID,
			Capabilities:  SupportedRunCapabilities(),
			Capacity:      protocol.RunCapacityReport{MaxJobs: cfg.MaxJobs},
		},
		journal: NewJobJournal(),
	}, nil
}

func (worker *Worker) Register(ctx context.Context) error {
	response, err := worker.client.Hello(ctx, protocol.RunHelloRequest{
		RegistrationToken: worker.cfg.RegistrationToken,
		RunEndpointID:     worker.cfg.RunEndpointID,
		DisplayName:       worker.cfg.DisplayName,
		Version:           worker.cfg.Version,
		Status:            "online",
		Platform:          runtime.GOOS,
		CapabilityReport: protocol.RunCapabilityReport{
			Capabilities: worker.state.Capabilities,
			Fingerprint:  capabilityFingerprint(worker.state.Capabilities),
		},
		Capacity: worker.capacityReport(),
	})
	if err != nil {
		return err
	}
	if !response.Accepted || response.SessionToken == "" {
		return fmt.Errorf("run hello was not accepted")
	}
	worker.state.SessionToken = response.SessionToken
	if sink, ok := worker.executor.logSink.(*SpoolLogSink); ok {
		sink.RunEndpointID = worker.state.RunEndpointID
		sink.SessionToken = worker.state.SessionToken
	}
	if hook, ok := worker.executor.artifactHook.(*QueueArtifactHook); ok {
		hook.RunEndpointID = worker.state.RunEndpointID
		hook.SessionToken = worker.state.SessionToken
	}
	return nil
}

func (worker *Worker) HeartbeatOnce(ctx context.Context) error {
	if worker.state.SessionToken == "" {
		return fmt.Errorf("worker is not registered")
	}
	response, err := worker.client.Heartbeat(ctx, protocol.RunHeartbeatRequest{
		RunEndpointID:         worker.state.RunEndpointID,
		SessionToken:          worker.state.SessionToken,
		Version:               worker.cfg.Version,
		Status:                "online",
		CapabilityFingerprint: capabilityFingerprint(worker.state.Capabilities),
		Capacity:              worker.capacityReport(),
	})
	if err != nil {
		return err
	}
	if !response.Accepted {
		return fmt.Errorf("heartbeat was not accepted")
	}
	worker.state.LastHeartbeat = response.ServerTime
	return nil
}

func (worker *Worker) ClaimAndRunOnce(ctx context.Context) (bool, error) {
	if worker.state.SessionToken == "" {
		return false, fmt.Errorf("worker is not registered")
	}
	claim, err := worker.client.ClaimJob(ctx, protocol.RunJobClaimRequest{
		RunEndpointID: worker.state.RunEndpointID,
		SessionToken:  worker.state.SessionToken,
		Capabilities:  worker.state.Capabilities,
		Capacity:      worker.capacityReport(),
	})
	if err != nil {
		return false, err
	}
	if !claim.Accepted || !claim.HasJob || claim.Job == nil {
		return false, nil
	}
	assignment := *claim.Job
	worker.journal.MarkActive(assignment)
	ack, err := worker.client.AckJob(ctx, protocol.RunJobAckRequest{
		RunEndpointID: worker.state.RunEndpointID,
		SessionToken:  worker.state.SessionToken,
		JobID:         assignment.JobID,
		LeaseToken:    assignment.LeaseToken,
		Attempt:       assignment.Attempt,
		Message:       "job accepted by run worker",
	})
	if err != nil {
		return true, err
	}
	assignment = ack.Job
	worker.journal.MarkActive(assignment)
	worker.state.Sequence++
	if _, err := worker.client.UpdateJobProgress(ctx, protocol.RunJobProgressRequest{
		RunEndpointID: worker.state.RunEndpointID,
		SessionToken:  worker.state.SessionToken,
		JobID:         assignment.JobID,
		LeaseToken:    assignment.LeaseToken,
		Attempt:       assignment.Attempt,
		Progress:      protocol.RunJobProgressReport{Percent: 10, Message: "lifecycle execution started"},
		Sequence:      worker.state.Sequence,
	}); err != nil {
		return true, err
	}
	jobCtx, cancel := context.WithCancel(ctx)
	cancelled := make(chan protocol.RunJobCancelPollResponse, 1)
	go func() {
		cancelPoll, pollErr := worker.client.PollJobCancel(ctx, protocol.RunJobCancelPollRequest{
			RunEndpointID: worker.state.RunEndpointID,
			SessionToken:  worker.state.SessionToken,
			JobID:         assignment.JobID,
			LeaseToken:    assignment.LeaseToken,
		})
		if pollErr == nil && cancelPoll.HasCancel {
			cancel()
			cancelled <- cancelPoll
			return
		}
		cancelled <- protocol.RunJobCancelPollResponse{Accepted: true}
	}()
	execution := worker.executor.ExecuteContext(jobCtx, assignment)
	cancel()
	select {
	case poll := <-cancelled:
		if poll.HasCancel && execution.State == lifecycleResultStateSucceeded {
			execution = LifecycleExecutionResult{
				State:     lifecycleResultStateCancelled,
				Progress:  protocol.RunJobProgressReport{Percent: 100, Message: "cancelled by platform"},
				Message:   "cancelled by platform",
				ErrorCode: "lifecycle_cancelled",
			}
		}
	default:
	}
	if _, err := worker.client.CompleteJob(ctx, LifecycleResultRequest(assignment, worker.state.SessionToken, execution)); err != nil {
		return true, err
	}
	worker.journal.MarkTerminal(assignment.JobID)
	return true, nil
}

func (worker *Worker) ReconcileOnce(ctx context.Context) error {
	if worker.state.SessionToken == "" {
		return fmt.Errorf("worker is not registered")
	}
	response, err := worker.client.ReconcileJobs(ctx, protocol.RunJobReconcileRequest{
		RunEndpointID: worker.state.RunEndpointID,
		SessionToken:  worker.state.SessionToken,
		ActiveJobIDs:  worker.journal.ActiveJobIDs(),
	})
	if err != nil {
		return err
	}
	for _, job := range response.ActiveJobs {
		worker.journal.MarkActive(job)
	}
	for _, unknown := range response.UnknownJobIDs {
		worker.journal.MarkTerminal(unknown)
	}
	return nil
}

func (worker *Worker) Run(ctx context.Context) error {
	if err := worker.Register(ctx); err != nil {
		return err
	}
	heartbeatInterval := durationOrDefault(worker.cfg.HeartbeatInterval, 15*time.Second)
	jobInterval := durationOrDefault(worker.cfg.PollInterval, 2*time.Second)
	heartbeatTicker := time.NewTicker(heartbeatInterval)
	jobTicker := time.NewTicker(jobInterval)
	defer heartbeatTicker.Stop()
	defer jobTicker.Stop()
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-heartbeatTicker.C:
			if err := worker.HeartbeatOnce(ctx); err != nil {
				heartbeatTicker.Reset(boundedRetryBackoff(worker.cfg.RetryBackoff))
				continue
			}
			heartbeatTicker.Reset(heartbeatInterval)
		case <-jobTicker.C:
			if _, err := worker.ClaimAndRunOnce(ctx); err != nil {
				jobTicker.Reset(boundedRetryBackoff(worker.cfg.RetryBackoff))
				continue
			}
			jobTicker.Reset(jobInterval)
		}
	}
}

func (worker *Worker) capacityReport() protocol.RunCapacityReport {
	return protocol.RunCapacityReport{
		MaxJobs:     worker.state.Capacity.MaxJobs,
		RunningJobs: worker.journal.ActiveCount(),
		QueuedJobs:  0,
		Summary:     "worker control active; job capacity reported separately",
	}
}

func (worker *Worker) State() WorkerState {
	state := worker.state
	state.Capabilities = append([]string(nil), state.Capabilities...)
	return state
}

type JobJournal struct {
	mu     sync.Mutex
	active map[string]protocol.RunJobAssignment
}

func NewJobJournal() *JobJournal {
	return &JobJournal{active: map[string]protocol.RunJobAssignment{}}
}

func (journal *JobJournal) MarkActive(job protocol.RunJobAssignment) {
	journal.mu.Lock()
	defer journal.mu.Unlock()
	journal.active[job.JobID] = job
}

func (journal *JobJournal) MarkTerminal(jobID string) {
	journal.mu.Lock()
	defer journal.mu.Unlock()
	delete(journal.active, jobID)
}

func (journal *JobJournal) ActiveJobIDs() []string {
	journal.mu.Lock()
	defer journal.mu.Unlock()
	ids := make([]string, 0, len(journal.active))
	for id := range journal.active {
		ids = append(ids, id)
	}
	return ids
}

func (journal *JobJournal) ActiveCount() int {
	journal.mu.Lock()
	defer journal.mu.Unlock()
	return len(journal.active)
}

type SpoolLogSink struct {
	RunEndpointID string
	SessionToken  string
	Spool         spool.LogSpool
	seq           uint64
}

func (sink *SpoolLogSink) Append(_ context.Context, assignment protocol.RunJobAssignment, stream string, line string) error {
	sink.seq++
	entry := protocol.LogEntry{Seq: sink.seq, Timestamp: time.Now().UTC(), Level: "info", Line: RedactText(line), Redacted: line != RedactText(line)}
	logStreamID := fmt.Sprintf("job.%s.%s", assignment.JobID, stream)
	return sink.Spool.Enqueue(protocol.LogBatchIngestRequest{
		RunEndpointID:    sink.RunEndpointID,
		SessionToken:     sink.SessionToken,
		LogStreamID:      logStreamID,
		ServerInstanceID: assignment.ServerInstanceID,
		StreamKey:        stream,
		Source:           "process",
		FirstSeq:         sink.seq,
		LastSeq:          sink.seq,
		Checksum:         checksumForText(entry.Line),
		Entries:          []protocol.LogEntry{entry},
	})
}

type QueueArtifactHook struct {
	RunEndpointID string
	SessionToken  string
	Queue         spool.ArtifactQueue
}

func (hook QueueArtifactHook) QueueLifecycleResult(_ context.Context, assignment protocol.RunJobAssignment, result ProcessResult) (string, error) {
	ref := fmt.Sprintf("artifact://jobs/%s/lifecycle-result", assignment.JobID)
	payload := []byte(RedactText(result.Stdout + result.Stderr))
	if len(payload) == 0 {
		payload = []byte("lifecycle result metadata")
	}
	artifactID := "artifact-" + assignment.JobID + "-lifecycle"
	if err := hook.Queue.Enqueue(protocol.ArtifactChunkUploadRequest{
		RunEndpointID: hook.RunEndpointID,
		SessionToken:  hook.SessionToken,
		TransferID:    "transfer-" + assignment.JobID,
		ArtifactID:    artifactID,
		ChunkIndex:    0,
		Offset:        0,
		SizeBytes:     len(payload),
		Checksum:      checksumForText(string(payload)),
		Payload:       payload,
	}); err != nil {
		return "", err
	}
	return ref, nil
}

func capabilityFingerprint(capabilities []string) string {
	return checksumForText(strings.Join(capabilities, ","))
}

func durationOrDefault(value time.Duration, fallback time.Duration) time.Duration {
	if value <= 0 {
		return fallback
	}
	return value
}

func boundedRetryBackoff(value time.Duration) time.Duration {
	value = durationOrDefault(value, time.Second)
	if value > 30*time.Second {
		return 30 * time.Second
	}
	return value
}