package redisactivity

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

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

const (
	ownershipAssign           = "assign"
	ownershipRenew            = "renew"
	ownershipBeginDrain       = "begin_drain"
	ownershipAcknowledgeDrain = "acknowledge_drain"
	ownershipGet              = "get"
	ownershipExpire           = "expire"
)

var _ ownershipDomain.Repository = (*Adapter)(nil)
var _ ownershipDomain.DrainTicketStore = (*Adapter)(nil)

func (a *Adapter) Assign(
	ctx context.Context,
	now time.Time,
	proxyID string,
	workerID string,
	ttl time.Duration,
) (ownershipDomain.Assignment, error) {
	if err := validateOwnershipCall(ctx, a); err != nil {
		return ownershipDomain.Assignment{}, err
	}
	if now.IsZero() || proxyID == "" || workerID == "" || ttl <= 0 {
		return ownershipDomain.Assignment{}, ownershipDomain.ErrInvalidOwnership
	}
	reply, err := a.runOwnership(ctx, ownershipAssign, true, now.UnixMilli(), proxyID, workerID, 0, durationMillis(ttl), 0, 0)
	if err != nil {
		return ownershipDomain.Assignment{}, err
	}
	switch reply.Status {
	case scriptOK:
		return decodeAssignmentReply(reply)
	case scriptAlreadyOwned:
		return ownershipDomain.Assignment{}, ownershipDomain.ErrAlreadyOwned
	case scriptUnavailable, scriptNotFound:
		return ownershipDomain.Assignment{}, ownershipDomain.ErrOwnershipUnavailable
	case scriptInvalid:
		return ownershipDomain.Assignment{}, ownershipDomain.ErrInvalidOwnership
	default:
		return ownershipDomain.Assignment{}, invalidScriptReply("unexpected ownership assign status")
	}
}

func (a *Adapter) Renew(
	ctx context.Context,
	now time.Time,
	proxyID string,
	workerID string,
	epoch uint64,
	ttl time.Duration,
) (ownershipDomain.Assignment, error) {
	if err := validateOwnershipCall(ctx, a); err != nil {
		return ownershipDomain.Assignment{}, err
	}
	if now.IsZero() || proxyID == "" || workerID == "" || epoch == 0 || ttl <= 0 {
		return ownershipDomain.Assignment{}, ownershipDomain.ErrInvalidOwnership
	}
	reply, err := a.runOwnership(ctx, ownershipRenew, true, now.UnixMilli(), proxyID, workerID, epoch, durationMillis(ttl), 0, 0)
	if err != nil {
		return ownershipDomain.Assignment{}, err
	}
	switch reply.Status {
	case scriptOK:
		return decodeAssignmentReply(reply)
	case scriptStale, scriptNotFound, scriptUnavailable:
		return ownershipDomain.Assignment{}, ownershipDomain.ErrStaleAssignment
	case scriptInvalid:
		return ownershipDomain.Assignment{}, ownershipDomain.ErrInvalidOwnership
	default:
		return ownershipDomain.Assignment{}, invalidScriptReply("unexpected ownership renew status")
	}
}

func (a *Adapter) BeginDrain(
	ctx context.Context,
	proxyID string,
	workerID string,
	epoch uint64,
) (ownershipDomain.Assignment, error) {
	if err := validateOwnershipCall(ctx, a); err != nil {
		return ownershipDomain.Assignment{}, err
	}
	if proxyID == "" || workerID == "" || epoch == 0 {
		return ownershipDomain.Assignment{}, ownershipDomain.ErrInvalidOwnership
	}
	reply, err := a.runOwnership(ctx, ownershipBeginDrain, true, 0, proxyID, workerID, epoch, 0, 0, 0)
	if err != nil {
		return ownershipDomain.Assignment{}, err
	}
	switch reply.Status {
	case scriptOK:
		return decodeAssignmentReply(reply)
	case scriptStale, scriptNotFound:
		return ownershipDomain.Assignment{}, ownershipDomain.ErrStaleAssignment
	case scriptInvalid:
		return ownershipDomain.Assignment{}, ownershipDomain.ErrInvalidOwnership
	default:
		return ownershipDomain.Assignment{}, invalidScriptReply("unexpected ownership drain status")
	}
}

func (a *Adapter) PendingDrains(ctx context.Context, workerID string, limit int) ([]ownershipDomain.DrainTicket, error) {
	if err := validateOwnershipCall(ctx, a); err != nil {
		return nil, err
	}
	if workerID == "" || limit <= 0 {
		return nil, ownershipDomain.ErrInvalidDrainTicket
	}
	result, err := runScript(ctx, a.client, drainTicketsScript, []string{
		a.keys.drainTickets, a.keys.workerDraining(workerID), a.keys.owners,
	}, workerID, limit)
	if err != nil {
		return nil, err
	}
	var reply drainTicketsScriptReply
	if err := decodeScriptResult(result, &reply); err != nil {
		return nil, err
	}
	if reply.Status == scriptInvalid {
		return nil, ownershipDomain.ErrInvalidDrainTicket
	}
	if reply.Status != scriptOK {
		return nil, invalidScriptReply("unexpected pending drain tickets reply")
	}
	tickets := make([]ownershipDomain.DrainTicket, 0, len(reply.Tickets))
	for _, raw := range reply.Tickets {
		var record drainTicketRecord
		if err := decodeJSON(raw, &record); err != nil || validateDrainTicketRecord(record) != nil {
			return nil, invalidScriptReply("pending drain tickets reply contained an invalid ticket")
		}
		ticket := ownershipDomain.DrainTicket{
			ProxyID: record.ProxyID, WorkerID: record.WorkerID, AssignmentEpoch: record.AssignmentEpoch,
			RequiredSnapshotEpoch: record.RequiredSnapshotEpoch,
		}
		if record.SnapshotVersion != 0 {
			checksum, err := hex.DecodeString(record.SnapshotChecksum)
			if err != nil || len(checksum) != sha256.Size {
				return nil, invalidScriptReply("pending drain ticket barrier checksum is invalid")
			}
			copy(ticket.Barrier.Checksum[:], checksum)
			ticket.Barrier = ownershipDomain.SnapshotBarrier{
				SessionID: record.SessionID, Version: record.SnapshotVersion,
				OwnershipEpoch: record.SnapshotOwnershipEpoch, Checksum: ticket.Barrier.Checksum,
			}
		}
		tickets = append(tickets, ticket)
	}
	return tickets, nil
}

func (a *Adapter) BindDrainBarrier(ctx context.Context, ticket ownershipDomain.DrainTicket) error {
	if err := validateOwnershipCall(ctx, a); err != nil {
		return err
	}
	if !validDrainTicket(ticket) || !validDrainBarrier(ticket.Barrier) ||
		ticket.Barrier.OwnershipEpoch < ticket.RequiredSnapshotEpoch {
		return ownershipDomain.ErrInvalidDrainTicket
	}
	result, err := runScript(ctx, a.client, bindDrainTicketScript, []string{
		a.keys.drainTickets, a.keys.owners, a.keys.workerDraining(ticket.WorkerID),
	}, ticket.ProxyID, ticket.WorkerID, ticket.AssignmentEpoch, ticket.RequiredSnapshotEpoch,
		ticket.Barrier.SessionID, ticket.Barrier.Version, ticket.Barrier.OwnershipEpoch, hex.EncodeToString(ticket.Barrier.Checksum[:]))
	if err != nil {
		return err
	}
	var reply ownershipScriptReply
	if err := decodeScriptResult(result, &reply); err != nil {
		return err
	}
	switch reply.Status {
	case scriptOK:
		return nil
	case scriptStale, scriptNotFound:
		return ownershipDomain.ErrStaleAssignment
	case scriptInvalid:
		return ownershipDomain.ErrInvalidDrainTicket
	default:
		return invalidScriptReply("unexpected drain ticket barrier reply")
	}
}

func (a *Adapter) AcknowledgeDrain(
	ctx context.Context,
	proxyID string,
	workerID string,
	epoch uint64,
	active int64,
	reserved int64,
) error {
	if err := validateOwnershipCall(ctx, a); err != nil {
		return err
	}
	if proxyID == "" || workerID == "" || epoch == 0 || active < 0 || reserved < 0 {
		return ownershipDomain.ErrInvalidOwnership
	}
	reply, err := a.runOwnership(ctx, ownershipAcknowledgeDrain, true, 0, proxyID, workerID, epoch, 0, active, reserved)
	if err != nil {
		return err
	}
	switch reply.Status {
	case scriptOK:
		return nil
	case scriptStale, scriptNotFound:
		return ownershipDomain.ErrStaleAssignment
	case scriptNotDraining:
		return ownershipDomain.ErrNotDraining
	case scriptDrainNotReady:
		return ownershipDomain.ErrDrainNotReady
	case scriptInvalid:
		return ownershipDomain.ErrInvalidOwnership
	default:
		return invalidScriptReply("unexpected ownership acknowledge status")
	}
}

func (a *Adapter) Get(ctx context.Context, proxyID string) (ownershipDomain.Assignment, bool, error) {
	if err := validateOwnershipCall(ctx, a); err != nil {
		return ownershipDomain.Assignment{}, false, err
	}
	if proxyID == "" {
		return ownershipDomain.Assignment{}, false, ownershipDomain.ErrInvalidOwnership
	}
	reply, err := a.runOwnership(ctx, ownershipGet, false, 0, proxyID, "", 0, 0, 0, 0)
	if err != nil {
		return ownershipDomain.Assignment{}, false, err
	}
	switch reply.Status {
	case scriptNotFound:
		return ownershipDomain.Assignment{}, false, nil
	case scriptOK:
		assignment, err := decodeAssignmentReply(reply)
		return assignment, err == nil, err
	case scriptInvalid:
		return ownershipDomain.Assignment{}, false, ownershipDomain.ErrInvalidOwnership
	default:
		return ownershipDomain.Assignment{}, false, invalidScriptReply("unexpected ownership get status")
	}
}

func (a *Adapter) Expire(ctx context.Context, now time.Time, limit int) ([]ownershipDomain.Assignment, error) {
	if err := validateOwnershipCall(ctx, a); err != nil {
		return nil, err
	}
	if now.IsZero() || limit <= 0 {
		return nil, ownershipDomain.ErrInvalidOwnership
	}
	reply, err := a.runOwnership(ctx, ownershipExpire, true, now.UnixMilli(), "", "", 0, int64(limit), 0, 0)
	if err != nil {
		return nil, err
	}
	if reply.Status == scriptInvalid {
		return nil, ownershipDomain.ErrInvalidOwnership
	}
	if reply.Status != scriptOK || reply.Record == "" {
		return nil, invalidScriptReply("unexpected ownership expire reply")
	}
	var records []ownershipRecord
	if err := decodeJSON(reply.Record, &records); err != nil {
		return nil, errors.Join(invalidScriptReply("ownership expire reply contained invalid records"), err)
	}
	assignments := make([]ownershipDomain.Assignment, 0, len(records))
	for _, record := range records {
		if err := validateOwnershipRecord(record); err != nil {
			return nil, errors.Join(invalidScriptReply("ownership expire reply contained invalid records"), err)
		}
		assignments = append(assignments, assignmentFromRecord(record))
	}
	if len(assignments) > limit {
		return nil, invalidScriptReply("ownership expire reply exceeded limit")
	}
	return assignments, nil
}

func (a *Adapter) runOwnership(
	ctx context.Context,
	operation string,
	mutating bool,
	nowMS int64,
	proxyID string,
	workerID string,
	epoch uint64,
	value int64,
	active int64,
	reserved int64,
) (ownershipScriptReply, error) {
	operationKey := a.keys.epoch
	if mutating {
		operationID, err := newOperationID()
		if err != nil {
			return ownershipScriptReply{}, err
		}
		operationKey = a.keys.operation(operationID)
	}
	result, err := runScript(ctx, a.client, ownershipScript, []string{
		a.keys.records, a.keys.unique, a.keys.idkeys, a.keys.expiry, a.keys.available,
		a.keys.inventory, a.keys.stateInventory, a.keys.owners, a.keys.ownerExpiry,
		a.keys.epoch, operationKey, a.keys.healthDue, a.keys.healthQueued, a.keys.healthLeases,
		a.keys.healthTasks, a.keys.healthTaskExpiry, a.keys.healthRefTask, a.keys.healthUnhealthy,
		a.keys.drainTickets,
	}, operation, operationTTLMillis(a.options.OperationTTL), a.options.CleanupLimit,
		nowMS, proxyID, workerID, epoch, value, active, reserved, a.keys.workerOwned(workerID), a.keys.workerDraining(workerID))
	if err != nil {
		return ownershipScriptReply{}, err
	}
	var reply ownershipScriptReply
	if err := decodeScriptResult(result, &reply); err != nil {
		return ownershipScriptReply{}, err
	}
	return reply, nil
}

func decodeAssignmentReply(reply ownershipScriptReply) (ownershipDomain.Assignment, error) {
	if reply.Record == "" {
		return ownershipDomain.Assignment{}, invalidScriptReply("ownership reply omitted assignment")
	}
	record, err := decodeOwnershipRecord(reply.Record)
	if err != nil {
		return ownershipDomain.Assignment{}, errors.Join(
			invalidScriptReply("ownership reply contained an invalid assignment"),
			fmt.Errorf("decode ownership assignment: %w", err),
		)
	}
	return assignmentFromRecord(record), nil
}

func assignmentFromRecord(record ownershipRecord) ownershipDomain.Assignment {
	return ownershipDomain.Assignment{
		ProxyID: record.ProxyID, WorkerID: record.WorkerID, Epoch: record.Epoch,
		Version: record.AssignmentVersion, ExpiresAt: time.UnixMilli(record.ExpiresAtMS).UTC(),
		Draining: record.Draining,
	}
}

func validateOwnershipCall(ctx context.Context, adapter *Adapter) error {
	if ctx == nil || adapter == nil {
		return ownershipDomain.ErrInvalidOwnership
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	return nil
}

func validDrainTicket(ticket ownershipDomain.DrainTicket) bool {
	return ticket.ProxyID != "" && ticket.WorkerID != "" && ticket.AssignmentEpoch > 0 &&
		ticket.RequiredSnapshotEpoch > ticket.AssignmentEpoch
}

func validDrainBarrier(barrier ownershipDomain.SnapshotBarrier) bool {
	return barrier.SessionID != "" && barrier.Version > 0 && barrier.OwnershipEpoch > 0 &&
		barrier.Checksum != [32]byte{}
}