package handler

import (
	"encoding/binary"
	"errors"
	"time"

	"ripple/bug"
	"ripple/commands"
	"ripple/crypto"
	"ripple/errmsgs"
	"ripple/helpers"
	"ripple/lockstep"
	"ripple/services"
	"ripple/state"
	"ripple/transport"
	"ripple/types"
)

var errInvalidStateTransition = errors.New("invalid state transition")

type LockstepHandlers struct {
	lockstep       *lockstep.Lockstep
	cptSender      *transport.CounterpartTransport
	resilience     *services.Resilience	
}

func NewLockstepHandlers(
	lockstep *lockstep.Lockstep,
	sender *transport.CounterpartTransport,
	resilience *services.Resilience,
) *LockstepHandlers {
	return &LockstepHandlers{
		lockstep:   lockstep,
		cptSender:  sender,
		resilience: resilience,
	}
}

func (h *LockstepHandlers) SetTrustlineInternal(st *state.State, userID types.UserIdentifier, args []byte) (func(), error) {
	if len(args) < 8 {
		return nil, errmsgs.ErrInvalidArgs
	}
	value := binary.BigEndian.Uint64(args[:8])
	account := st.Storage.MustGetAccount(userID)
	account.TrustlineOut = value
	st.Storage.Accounts[userID] = account
	return nil, nil
}

func (h *LockstepHandlers) SetTrustlineExternal(st *state.State, userID types.UserIdentifier, args []byte) (func(), error) {
	if len(args) < 8 {
		return nil, errmsgs.ErrInvalidArgs
	}
	value := binary.BigEndian.Uint64(args[:8])
	account := st.Storage.MustGetAccount(userID)
	account.TrustlineIn = value
	st.Storage.Accounts[userID] = account
	return nil, nil
}

func (h *LockstepHandlers) InitPaymentInternal(st *state.State, userID types.UserIdentifier, args []byte) (func(), error) {
	if len(args) < 60 {
		return nil, errmsgs.ErrInvalidArgs
	}
	var paymentID [32]byte
	copy(paymentID[:], args[:32])
	if crypto.Sha256(st.Memory.Preimage[:]) != paymentID {
		return nil, nil
	}
	st.Storage.Preimage = st.Memory.Preimage
	return func() {
		h.lockstep.Enqueue(userID, types.Instruction{
			Command:   commands.LOCKSTEP_COMMIT_PAYMENT,
			Arguments: args,
		})
	}, nil
}

func (h *LockstepHandlers) InitPaymentExternal(st *state.State, userID types.UserIdentifier, args []byte) (func(), error) {
	if len(args) < 60 {
		return nil, errmsgs.ErrInvalidArgs
	}
	return nil, nil
}

func (h *LockstepHandlers) CommitPaymentInternal(st *state.State, userID types.UserIdentifier, args []byte) (func(), error) {
	if len(args) < 60 {
		return nil, errmsgs.ErrInvalidArgs
	}
	var paymentID [32]byte
	copy(paymentID[:], args[:32])
	amount      := binary.BigEndian.Uint64(args[32:40])
	penaltyRate := binary.BigEndian.Uint32(args[40:44])
	feeOut      := binary.BigEndian.Uint64(args[44:52])
	tax         := binary.BigEndian.Uint64(args[52:60])

	if st.Storage.GetBandwidthOut(userID) < int64(amount+feeOut+tax) {
		return nil, errInvalidStateTransition
	}
	{
		account := st.Storage.MustGetAccount(userID)
		account.Pending[paymentID] = struct{}{}
		st.Storage.Accounts[userID] = account
	}
	payment, ok := st.Storage.Payments[paymentID]
	if !ok {
		newPayment := state.Payment{
			Amount:      amount,
			Outgoing:    userID,
			PenaltyRate: penaltyRate,
			FeeOut:      feeOut,
			Tax:         tax,
			Preimage:    st.Storage.Preimage,
			CreatedAt:   time.Now().Unix(),
		}
		pf, ok := st.Memory.GetPathfinding(paymentID)
		if !ok {
			newPayment.Status = state.Cancel
			st.Storage.Payments[paymentID] = newPayment
			return func() {
				h.lockstep.Enqueue(userID, helpers.BuildCancelIncoming(paymentID, newPayment))
			}, nil
		}
		if int64(amount) != helpers.Abs(pf.Amount) || penaltyRate != pf.PenaltyRate || feeOut != uint64(pf.Hops)*pf.Fee || tax != pf.Tax {
			panic(bug.BugStateViolated)
		}
		delete(st.Memory.Pathfinding, paymentID)
		newPayment.Counterpart = pf.Counterpart
		st.Storage.Payments[paymentID] = newPayment
		return nil, nil
	}
	payment.Outgoing = userID
	payment.FeeOut = feeOut
	st.Storage.Payments[paymentID] = payment
	return nil, nil
}

func (h *LockstepHandlers) CommitPaymentExternal(st *state.State, userID types.UserIdentifier, args []byte) (func(), error) {
	if len(args) < 60 {
		return nil, errmsgs.ErrInvalidArgs
	}
	var paymentID [32]byte
	copy(paymentID[:], args[:32])
	amount      := binary.BigEndian.Uint64(args[32:40])
	penaltyRate := binary.BigEndian.Uint32(args[40:44])
	feeIn       := binary.BigEndian.Uint64(args[44:52])
	tax         := binary.BigEndian.Uint64(args[52:60])

	if st.Storage.GetBandwidthIn(userID) < int64(amount+feeIn+tax) {
		return nil, errInvalidStateTransition
	}
	{
		account := st.Storage.MustGetAccount(userID)
		account.Pending[paymentID] = struct{}{}
		st.Storage.Accounts[userID] = account
	}

	payment := state.Payment{
		Amount:      amount,
		Incoming:    userID,
		PenaltyRate: penaltyRate,
		FeeIn:       feeIn,
		Tax:         tax,
		CreatedAt:   time.Now().Unix(),
	}

	pf, ok := st.Memory.GetPathfinding(paymentID)
	if !ok || int64(amount) != helpers.Abs(pf.Amount) || penaltyRate != pf.PenaltyRate || feeIn != uint64(pf.Hops)*pf.Fee+pf.Fee || tax != pf.Tax {
		payment.Status = state.Cancel
		st.Storage.Payments[paymentID] = payment
		return func() {
			h.lockstep.Enqueue(userID, helpers.BuildCancelIncoming(paymentID, payment))
		}, nil
	}
	
	delete(st.Memory.Pathfinding, paymentID)
	if pf.Counterpart != (types.UserIdentifier{}) {
		payment.Counterpart = pf.Counterpart
		payment.Preimage = st.Memory.Preimage
		st.Storage.Payments[paymentID] = payment
		return func() {
			h.cptSender.Send(types.Instruction{Command: commands.COUNTERPART_SEAL_PAYMENT})
		}, nil
	}

	st.Storage.Payments[paymentID] = payment

	return func() {
		h.lockstep.Enqueue(pf.Outgoing, helpers.BuildCommitPayment(paymentID, amount, penaltyRate, feeIn-pf.Fee, tax))
	}, nil
}

func (h *LockstepHandlers) SealPaymentInternal(st *state.State, userID types.UserIdentifier, args []byte) (func(), error) {
	if len(args) < 40 {
		return nil, errmsgs.ErrInvalidArgs
	}
	var paymentID [32]byte
	copy(paymentID[:], args[:32])
	commitPenalty := binary.BigEndian.Uint64(args[32:40])

	payment, ok := st.Storage.Payments[paymentID]
	if !ok {
		panic(bug.BugStateViolated)
	}
	if !payment.Committed() && !payment.SealIn() || payment.Outgoing != userID {
		return nil, errInvalidStateTransition
	}
	payment.CommitPenalty = commitPenalty

	if commitPenalty >= payment.Amount {
		payment.Outgoing = types.UserIdentifier{}
		amount := payment.Amount+payment.FeeOut
		st.Storage.MustFinalizeOutgoing(userID, paymentID, amount+payment.Tax, payment.GetTaxRate(amount))
		if payment.Incoming == (types.UserIdentifier{}) {
			delete(st.Storage.Payments, paymentID)
			return nil, nil
		}
	}

	payment.Status = state.Sealed
	st.Storage.Payments[paymentID] = payment
	return nil, nil
}

func (h *LockstepHandlers) SealPaymentExternal(st *state.State, userID types.UserIdentifier, args []byte) (func(), error) {
	if len(args) < 40 {
		return nil, errmsgs.ErrInvalidArgs
	}
	var paymentID [32]byte
	copy(paymentID[:], args[:32])
	commitPenalty := binary.BigEndian.Uint64(args[32:40])

	payment, ok := st.Storage.Payments[paymentID]
	if !ok || !payment.Committed() || payment.Incoming != userID {
		return nil, errInvalidStateTransition
	}
	payment.CommitPenalty = commitPenalty

	var redistrib func()

	if commitPenalty >= payment.Amount {
		payment.Incoming = types.UserIdentifier{}
		amount := payment.Amount+payment.FeeIn
		taxrate := payment.GetTaxRate(amount)
		st.MustFinalizeIncoming(userID, paymentID, amount+payment.Tax, taxrate)
        redistrib = func() {
            h.resilience.Redistribute(userID, payment.Tax, taxrate)
        }
		if payment.Outgoing == (types.UserIdentifier{}) {
			delete(st.Storage.Payments, paymentID)
			return redistrib, nil
		}
	}
	if payment.Outgoing == (types.UserIdentifier{}) {
		payment.Status = state.Sealed
		st.Storage.Payments[paymentID] = payment
		return func() {
			h.lockstep.Enqueue(payment.Incoming, helpers.BuildFinalizePayment(paymentID, payment))
			if redistrib != nil {
				redistrib()
			}
		}, nil
	} else {
		payment.Status = state.SealIn
		st.Storage.Payments[paymentID] = payment
		return func() {
			h.lockstep.Enqueue(payment.Outgoing, helpers.BuildSealPayment(paymentID, payment))
			if redistrib != nil {
				redistrib()
			}
		}, nil
	}
}

func (h *LockstepHandlers) FinalizePaymentInternal(st *state.State, userID types.UserIdentifier, args []byte) (func(), error) {
	if len(args) < 72 {
		return nil, errmsgs.ErrInvalidArgs
	}
	var paymentID [32]byte
	copy(paymentID[:], args[:32])
	var preimage [32]byte
	copy(preimage[:], args[32:64])
	amount := binary.BigEndian.Uint64(args[64:72])

	if crypto.Sha256(preimage[:]) != paymentID {
		return nil, errInvalidStateTransition
	}

	payment, ok := st.Storage.Payments[paymentID]
	if !ok {
		panic(bug.BugStateViolated)
	}

	if !payment.Sealed() || payment.Incoming != userID || amount > payment.Amount+payment.FeeIn {
		return nil, errInvalidStateTransition
	}
	taxrate := payment.GetTaxRate(amount)
	st.MustFinalizeIncoming(userID, paymentID, amount+payment.Tax, taxrate)

	if payment.Counterpart != (types.UserIdentifier{}) {
		st.Storage.AddReceipt(state.Receipt{
			Identifier:  paymentID,
			Counterpart: payment.Counterpart,
			Amount:      int64(payment.Amount),
			Timestamp:   time.Now().Unix(),
		})
	}

	delete(st.Storage.Payments, paymentID)
	return func() {
		h.resilience.Redistribute(userID, payment.Tax, taxrate)
	}, nil
}

func (h *LockstepHandlers) FinalizePaymentExternal(st *state.State, userID types.UserIdentifier, args []byte) (func(), error) {
	if len(args) < 72 {
		return nil, errmsgs.ErrInvalidArgs
	}
	var paymentID [32]byte
	copy(paymentID[:], args[:32])
	var preimage [32]byte
	copy(preimage[:], args[32:64])
	amount := binary.BigEndian.Uint64(args[64:72])

	if crypto.Sha256(preimage[:]) != paymentID {
		return nil, errInvalidStateTransition
	}
	
	payment, ok := st.Storage.Payments[paymentID]
	if !ok || !payment.Sealed() || payment.Outgoing != userID || amount > payment.Amount+payment.FeeOut {
		return nil, errInvalidStateTransition
	}
	st.Storage.MustFinalizeOutgoing(userID, paymentID, amount+payment.Tax, payment.GetTaxRate(amount))

	if payment.Counterpart != (types.UserIdentifier{}) {
		st.Storage.AddReceipt(state.Receipt{
			Identifier:  paymentID,
			Counterpart: payment.Counterpart,
			Amount:      -int64(payment.Amount),
			Timestamp:   time.Now().Unix(),
		})
		delete(st.Storage.Payments, paymentID)
		return nil, nil
	}
	payment.Preimage = preimage
	payment.Outgoing = types.UserIdentifier{}
	st.Storage.Payments[paymentID] = payment

	return func() {
		h.lockstep.Enqueue(payment.Incoming, helpers.BuildFinalizePayment(paymentID, payment))
	}, nil
}

func (h *LockstepHandlers) CancelPaymentInternal(st *state.State, userID types.UserIdentifier, args []byte) (func(), error) {
	if len(args) < 72 {
		return nil, errmsgs.ErrInvalidArgs
	}
	var paymentID [32]byte
	copy(paymentID[:], args[:32])
	var preimage [32]byte
	copy(preimage[:], args[32:64])
	amount := binary.BigEndian.Uint64(args[64:72])

	payment, ok := st.Storage.Payments[paymentID]
	if !ok || !payment.Cancelled() {
		panic(bug.BugStateViolated)
	}

	if payment.Outgoing == userID {
		if crypto.Sha256(preimage[:]) != paymentID {
			return nil, errInvalidStateTransition
		}
		st.Storage.MustFinalizeOutgoing(userID, paymentID, amount, 0)
	} else if payment.Incoming == userID {
		st.MustFinalizeIncoming(userID, paymentID, amount, 0)
	} else {
		return nil, errInvalidStateTransition
	}

	delete(st.Storage.Payments, paymentID)
	return nil, nil
}

func (h *LockstepHandlers) CancelPaymentExternal(st *state.State, userID types.UserIdentifier, args []byte) (func(), error) {
	if len(args) < 72 {
		return nil, errmsgs.ErrInvalidArgs
	}
	var paymentID [32]byte
	copy(paymentID[:], args[:32])
	var preimage [32]byte
	copy(preimage[:], args[32:64])
	amount := binary.BigEndian.Uint64(args[64:72])

	payment, ok := st.Storage.Payments[paymentID]
	if !ok || payment.SealIn() || payment.Sealed() {
		return nil, errInvalidStateTransition
	}

	if payment.Incoming == userID {
		if crypto.Sha256(preimage[:]) != paymentID {
			return nil, errInvalidStateTransition
		}
		payment.Preimage = preimage
		st.MustFinalizeIncoming(userID, paymentID, amount, 0)
	} else if payment.Outgoing == userID {
		st.Storage.MustFinalizeOutgoing(userID, paymentID, amount, 0)
	} else {
		return nil, errInvalidStateTransition
	}

	if payment.Incoming == userID {
		if payment.Outgoing == (types.UserIdentifier{}) {
			delete(st.Storage.Payments, paymentID)
			return nil, nil
		}
		payment.Incoming = types.UserIdentifier{}
	} else {
		if payment.Incoming == (types.UserIdentifier{}) {
			delete(st.Storage.Payments, paymentID)
			return nil, nil
		}
		payment.Outgoing = types.UserIdentifier{}
	}
	payment.Status = state.Cancel
	st.Storage.Payments[paymentID] = payment

	ok, nextHop, instr := helpers.BuildCancelAuto(paymentID, payment)
	if ok {
		return func() {
			h.lockstep.Enqueue(nextHop, instr)
		}, nil
	}
	return nil, nil
}

func (h *LockstepHandlers) CleanupPaymentInternal(st *state.State, userID types.UserIdentifier, args []byte) (func(), error) {
	if len(args) < 32 {
		return nil, errmsgs.ErrInvalidArgs
	}
	var paymentID [32]byte
	copy(paymentID[:], args[:32])

	payment, ok := st.Storage.Payments[paymentID]
	if !ok {
		return nil, errInvalidStateTransition
	}

	if payment.Committed() {
		payment.CommitPenalty = payment.Amount
	}

	switch {
	case payment.Incoming == userID:
		amount := payment.CommitPenalty+payment.FeeIn
		taxrate := payment.GetTaxRate(amount)
		st.MustFinalizeIncoming(userID, paymentID, amount+payment.Tax, taxrate)
		payment.Incoming = types.UserIdentifier{}
		st.Storage.Payments[paymentID] = payment
		if payment.Outgoing == (types.UserIdentifier{}) {
			delete(st.Storage.Payments, paymentID)
		}
		return func() {
			h.resilience.Redistribute(userID, payment.Tax, taxrate)
		}, nil

	case payment.Outgoing == userID:
		amount := payment.CommitPenalty+payment.FeeOut
		st.Storage.MustFinalizeOutgoing(userID, paymentID, amount+payment.Tax, payment.GetTaxRate(amount))
		payment.Outgoing = types.UserIdentifier{}
		st.Storage.Payments[paymentID] = payment
		if payment.Incoming == (types.UserIdentifier{}) {
			delete(st.Storage.Payments, paymentID)
		}
		return nil, nil

	default:
		return nil, errInvalidStateTransition
	}
}

func (h *LockstepHandlers) CleanupPaymentExternal(st *state.State, userID types.UserIdentifier, args []byte) (func(), error) {
	if len(args) < 32 {
		return nil, errmsgs.ErrInvalidArgs
	}
	var paymentID [32]byte
	copy(paymentID[:], args[:32])

	payment, ok := st.Storage.Payments[paymentID]
	if !ok {
		return nil, errInvalidStateTransition
	}

	if payment.Committed() {
		payment.CommitPenalty = payment.Amount
	}
	
	switch {
	case payment.Incoming == userID:
		amount := payment.CommitPenalty+payment.FeeIn
		taxrate := payment.GetTaxRate(amount)
		st.MustFinalizeIncoming(userID, paymentID, amount+payment.Tax, taxrate)
		payment.Incoming = types.UserIdentifier{}
		st.Storage.Payments[paymentID] = payment
		if payment.Outgoing == (types.UserIdentifier{}) {
			delete(st.Storage.Payments, paymentID)
		} else if payment.PenaltyTicker() >= payment.Amount {
			return func() {
				h.lockstep.Enqueue(payment.Outgoing, helpers.BuildCleanupPayment(paymentID))
				h.resilience.Redistribute(userID, payment.Tax, taxrate)
			}, nil
		}
		return func() {
			h.resilience.Redistribute(userID, payment.Tax, taxrate)
		}, nil

	case payment.Outgoing == userID:
		amount := payment.CommitPenalty+payment.FeeOut
		st.Storage.MustFinalizeOutgoing(userID, paymentID, amount+payment.Tax, payment.GetTaxRate(amount))
		payment.Outgoing = types.UserIdentifier{}
		st.Storage.Payments[paymentID] = payment
		if payment.Incoming == (types.UserIdentifier{}) {
			delete(st.Storage.Payments, paymentID)
		} else if payment.PenaltyTicker() >= payment.Amount {
			return func() {
				h.lockstep.Enqueue(payment.Incoming, helpers.BuildCleanupPayment(paymentID))
			}, nil
		}
		return nil, nil
	default:
		return nil, errInvalidStateTransition
	}
}

func (h *LockstepHandlers) SwarmRedistributionInternal(st *state.State, userID types.UserIdentifier, args []byte) (func(), error) {
	if len(args) < 8 {
		return nil, errmsgs.ErrInvalidArgs
	}
	amount := int64(binary.BigEndian.Uint64(args[:8]))
	acc := st.Storage.MustGetAccount(userID)
	if amount > acc.Creditline {
		return nil, errmsgs.ErrInvalidArgs
	}
	acc.Creditline -= amount
	st.Storage.Accounts[userID] = acc
	return nil, nil
}

func (h *LockstepHandlers) SwarmRedistributionExternal(st *state.State, userID types.UserIdentifier, args []byte) (func(), error) {
	if len(args) < 8 {
		return nil, errmsgs.ErrInvalidArgs
	}
	amount := binary.BigEndian.Uint64(args[:8])
	acc := st.Storage.MustGetAccount(userID)
	if int64(amount) > -acc.Creditline {
		return nil, errmsgs.ErrInvalidArgs
	}
	acc.Creditline += int64(amount)
	st.Storage.Accounts[userID] = acc

	return func() {
		h.resilience.Redistribute(userID, amount, acc.Width)
	}, nil
}