package worker

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"time"

	ownershipDomain "proxy-pool/internal/domain/ownership"
	"proxy-pool/internal/domain/workerruntime"
)

var (
	ErrInvalidCommand  = errors.New("invalid worker control command")
	ErrProtocolVersion = errors.New("unsupported worker protocol version")
	ErrUnavailable     = errors.New("worker control service unavailable")
)

type RegisterCommand struct {
	WorkerID        string
	InstanceID      string
	Zone            string
	ProtocolVersion uint32
	Labels          map[string]string
}

type SnapshotAcknowledgement struct {
	WorkerID       string
	SessionID      string
	Version        uint64
	OwnershipEpoch uint64
	Checksum       []byte
	Applied        bool
	ErrorCode      string
	ErrorMessage   string
}

type Registration struct {
	WorkerID          string
	SessionID         string
	OwnershipEpoch    uint64
	HeartbeatInterval time.Duration
	MaxStaleAge       time.Duration
}

type RuntimeDecision struct {
	AcceptedOwnershipEpoch uint64
	RequireFullSnapshot    bool
}

type Options struct {
	ProtocolVersion    uint32
	HeartbeatInterval  time.Duration
	SessionTTL         time.Duration
	MaxStaleAge        time.Duration
	MaxRuntimeCounters int
	MaxSnapshotBytes   int
	SnapshotReader     ownershipDomain.SnapshotReader
	SessionID          func() (string, error)
}

type Service interface {
	Register(context.Context, RegisterCommand) (Registration, error)
	CurrentOwnershipEpoch(context.Context) (uint64, error)
	ValidateSession(context.Context, string, string) error
	IssueSnapshot(context.Context, string, workerruntime.SnapshotReference) error
	Acknowledge(context.Context, SnapshotAcknowledgement) error
	ReportRuntime(context.Context, workerruntime.Report) (RuntimeDecision, error)
}

func (service *service) CurrentOwnershipEpoch(ctx context.Context) (uint64, error) {
	if ctx == nil {
		return 0, ErrInvalidCommand
	}
	if err := ctx.Err(); err != nil {
		return 0, err
	}
	epoch, err := service.store.CurrentOwnershipEpoch(ctx)
	if err != nil {
		return 0, classifyStoreError(err)
	}
	return epoch, nil
}

func (service *service) ValidateSession(ctx context.Context, workerID, sessionID string) error {
	if ctx == nil || !workerruntime.ValidIdentifier(workerID) || !workerruntime.ValidIdentifier(sessionID) {
		return ErrInvalidCommand
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	if err := service.store.ValidateSession(ctx, workerID, sessionID); err != nil {
		return classifyStoreError(err)
	}
	return nil
}

func (service *service) IssueSnapshot(ctx context.Context, sessionID string, reference workerruntime.SnapshotReference) error {
	if ctx == nil {
		return ErrInvalidCommand
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	normalized, err := workerruntime.NormalizeSnapshotReference(reference)
	if err != nil || !workerruntime.ValidIdentifier(sessionID) {
		return errors.Join(ErrInvalidCommand, err)
	}
	if err := service.store.RecordIssuedSnapshot(ctx, sessionID, normalized, service.options.SessionTTL); err != nil {
		return classifyStoreError(err)
	}
	return nil
}

type service struct {
	store     workerruntime.ControlStore
	options   Options
	snapshots SnapshotSource
}

func NewService(store workerruntime.ControlStore, options Options) (Service, error) {
	if store == nil || options.ProtocolVersion == 0 || options.HeartbeatInterval <= 0 ||
		options.SessionTTL < 3*options.HeartbeatInterval || options.MaxStaleAge < options.HeartbeatInterval ||
		options.MaxRuntimeCounters <= 0 {
		return nil, ErrInvalidCommand
	}
	if options.SessionID == nil {
		options.SessionID = randomSessionID
	}
	result := &service{store: store, options: options}
	if options.SnapshotReader != nil {
		source, err := NewOwnedSnapshotSource(
			result, options.SnapshotReader, options.MaxStaleAge, options.MaxRuntimeCounters, options.MaxSnapshotBytes, time.Now,
		)
		if err != nil {
			return nil, errors.Join(ErrInvalidCommand, err)
		}
		result.snapshots = source
	}
	return result, nil
}

func (service *service) SnapshotSource() SnapshotSource {
	if service == nil {
		return nil
	}
	return service.snapshots
}

func (service *service) Register(ctx context.Context, command RegisterCommand) (Registration, error) {
	if ctx == nil {
		return Registration{}, ErrInvalidCommand
	}
	if err := ctx.Err(); err != nil {
		return Registration{}, err
	}
	if command.ProtocolVersion != service.options.ProtocolVersion {
		return Registration{}, ErrProtocolVersion
	}
	epoch, err := service.store.CurrentOwnershipEpoch(ctx)
	if err != nil {
		return Registration{}, classifyStoreError(err)
	}
	sessionID, err := service.options.SessionID()
	if err != nil || !workerruntime.ValidIdentifier(sessionID) {
		return Registration{}, errors.Join(ErrUnavailable, err)
	}
	session := workerruntime.Session{
		WorkerID: command.WorkerID, InstanceID: command.InstanceID, SessionID: sessionID,
		Zone: command.Zone, ProtocolVersion: command.ProtocolVersion, Labels: command.Labels,
	}
	if _, err := workerruntime.NormalizeSession(session); err != nil {
		return Registration{}, errors.Join(ErrInvalidCommand, err)
	}
	if err := service.store.OpenSession(ctx, session, service.options.SessionTTL); err != nil {
		return Registration{}, classifyStoreError(err)
	}
	return Registration{
		WorkerID: command.WorkerID, SessionID: sessionID, OwnershipEpoch: epoch,
		HeartbeatInterval: service.options.HeartbeatInterval, MaxStaleAge: service.options.MaxStaleAge,
	}, nil
}

func (service *service) Acknowledge(ctx context.Context, acknowledgement SnapshotAcknowledgement) error {
	if ctx == nil || len(acknowledgement.Checksum) != sha256.Size || len(acknowledgement.ErrorMessage) > 512 {
		return ErrInvalidCommand
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	var checksum [sha256.Size]byte
	copy(checksum[:], acknowledgement.Checksum)
	domainAcknowledgement := workerruntime.SnapshotAcknowledgement{
		WorkerID: acknowledgement.WorkerID, SessionID: acknowledgement.SessionID,
		Reference: workerruntime.SnapshotReference{
			WorkerID: acknowledgement.WorkerID, Version: acknowledgement.Version,
			OwnershipEpoch: acknowledgement.OwnershipEpoch, Checksum: checksum,
		},
		Applied: acknowledgement.Applied, ErrorCode: acknowledgement.ErrorCode,
	}
	if _, err := workerruntime.NormalizeAcknowledgement(domainAcknowledgement); err != nil {
		return errors.Join(ErrInvalidCommand, err)
	}
	if err := service.store.AcknowledgeSnapshot(ctx, domainAcknowledgement, service.options.SessionTTL); err != nil {
		return classifyStoreError(err)
	}
	return nil
}

func (service *service) ReportRuntime(ctx context.Context, report workerruntime.Report) (RuntimeDecision, error) {
	if ctx == nil {
		return RuntimeDecision{}, ErrInvalidCommand
	}
	if err := ctx.Err(); err != nil {
		return RuntimeDecision{}, err
	}
	normalized, _, err := workerruntime.NormalizeReport(report)
	if err != nil || len(normalized.Counters) > service.options.MaxRuntimeCounters {
		return RuntimeDecision{}, errors.Join(ErrInvalidCommand, err)
	}
	if err := service.store.ReplaceRuntime(ctx, normalized, service.options.SessionTTL); err != nil {
		if errors.Is(err, workerruntime.ErrSnapshotMismatch) {
			epoch, epochErr := service.store.CurrentOwnershipEpoch(ctx)
			if epochErr != nil {
				return RuntimeDecision{}, classifyStoreError(epochErr)
			}
			return RuntimeDecision{AcceptedOwnershipEpoch: epoch, RequireFullSnapshot: true}, nil
		}
		return RuntimeDecision{}, classifyStoreError(err)
	}
	epoch, err := service.store.CurrentOwnershipEpoch(ctx)
	if err != nil {
		return RuntimeDecision{}, classifyStoreError(err)
	}
	return RuntimeDecision{AcceptedOwnershipEpoch: epoch}, nil
}

func randomSessionID() (string, error) {
	var value [16]byte
	if _, err := rand.Read(value[:]); err != nil {
		return "", err
	}
	return hex.EncodeToString(value[:]), nil
}

func classifyStoreError(err error) error {
	if err == nil {
		return nil
	}
	if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) ||
		errors.Is(err, workerruntime.ErrInvalidSession) || errors.Is(err, workerruntime.ErrInvalidReport) ||
		errors.Is(err, workerruntime.ErrInvalidAcknowledgement) || errors.Is(err, workerruntime.ErrStaleSession) ||
		errors.Is(err, workerruntime.ErrStaleReport) || errors.Is(err, workerruntime.ErrConflictingReport) ||
		errors.Is(err, workerruntime.ErrSnapshotMismatch) || errors.Is(err, workerruntime.ErrStaleAcknowledgement) {
		return err
	}
	return errors.Join(ErrUnavailable, err)
}