package schema

import (
	"context"
	"fmt"

	"entgo.io/ent"
	"entgo.io/ent/dialect/entsql"
	"entgo.io/ent/schema/edge"
	"entgo.io/ent/schema/field"
	"entgo.io/ent/schema/index"
	"github.com/lightsparkdev/spark/common"
	"github.com/lightsparkdev/spark/common/keys"
	gen "github.com/lightsparkdev/spark/so/ent"
	"github.com/lightsparkdev/spark/so/ent/hook"
	st "github.com/lightsparkdev/spark/so/ent/schema/schematype"
)

// TreeNode is the schema for the tree nodes table.
type TreeNode struct {
	ent.Schema
}

// Mixin is the mixin for the tree nodes table.
func (TreeNode) Mixin() []ent.Mixin {
	return []ent.Mixin{
		BaseMixin{},
	}
}

// Fields are the fields for the tree nodes table.
func (TreeNode) Fields() []ent.Field {
	return []ent.Field{
		field.Uint64("value").Immutable(),
		field.Enum("status").GoType(st.TreeNodeStatus("")),
		field.Bytes("verifying_pubkey").Immutable().GoType(keys.Public{}),
		field.Bytes("owner_identity_pubkey").GoType(keys.Public{}),
		field.Bytes("owner_signing_pubkey").GoType(keys.Public{}),
		field.Int16("vout"),

		field.Uint64("node_confirmation_height").Optional(),
		field.Uint64("refund_confirmation_height").Optional(),

		// Node transactions
		field.Bytes("raw_tx").NotEmpty(),
		field.Bytes("direct_tx").Optional(),
		field.Bytes("direct_from_cpfp_refund_tx").Optional(),
		field.Bytes("raw_txid").Optional().Comment("Valid transaction ID of the stored node transaction. Generated by hooks from raw_tx."),
		field.Bytes("direct_txid").Optional().Comment("Valid transaction ID of the stored direct node transaction. Generated by hooks from direct_tx."),
		field.Bytes("direct_from_cpfp_refund_txid").Optional().Comment("Valid transaction ID of the stored direct from CPFP node transaction. Generated by hooks from direct_from_cpfp_refund_tx."),

		// Refund transactions
		field.Bytes("raw_refund_tx").Optional().Comment("A transaction to exit Spark unilaterally. Only leafs have this transaction."),
		field.Bytes("direct_refund_tx").Optional(),
		field.Bytes("raw_refund_txid").Optional().Comment("Valid transaction ID of the stored refund transaction. Generated by hooks from raw_refund_tx."),
		field.Bytes("direct_refund_txid").Optional().Comment("Valid transaction ID of the direct refund transaction. Generated by hooks from direct_refund_tx."),
	}
}

// Edges are the edges for the tree nodes table.
func (TreeNode) Edges() []ent.Edge {
	return []ent.Edge{
		edge.To("tree", Tree.Type).
			Unique().
			Required(),
		edge.To("parent", TreeNode.Type).
			Unique(),
		edge.To("signing_keyshare", SigningKeyshare.Type).
			Unique().
			Required(),
		edge.From("children", TreeNode.Type).Ref("parent"),
	}
}

// Indexes are the indexes for the tree nodes table.
func (TreeNode) Indexes() []ent.Index {
	return []ent.Index{
		index.Edges("parent"),
		index.Edges("tree"),
		index.Edges("signing_keyshare"),
		index.Fields("owner_identity_pubkey"),
		index.Fields("owner_identity_pubkey", "status"),
		index.Fields("node_confirmation_height"),
		index.Fields("refund_confirmation_height"),
		index.Fields("update_time"),

		index.Fields("raw_txid").Annotations(
			entsql.IndexWhere("raw_txid is not null"),
		),
		index.Fields("direct_txid").Annotations(
			entsql.IndexWhere("direct_txid is not null"),
		),
		index.Fields("direct_from_cpfp_refund_txid").Annotations(
			entsql.IndexWhere("direct_from_cpfp_refund_txid is not null"),
		),

		index.Fields("raw_refund_txid").Annotations(
			entsql.IndexWhere("raw_refund_txid is not null"),
		),
		index.Fields("direct_refund_txid").Annotations(
			entsql.IndexWhere("direct_refund_txid is not null"),
		),
	}
}

func (TreeNode) Hooks() []ent.Hook {
	return []ent.Hook{
		// Define helper hooks that will populate TXIDs for the tranactions in the mutation.
		// Clearing a field will also clear a corresponding TXID.
		// If a transaction is set to nil the corresponding TXID will be cleared.
		// Be midful that updates on transactions are slower because the DB indexes on TXIDs are also updated.
		// The first hook handles single entity updates and the second hook handles creations.
		// These hooks will not be called if the entity is part of a batch update.
		hook.On(
			func(next ent.Mutator) ent.Mutator {
				return hook.TreeNodeFunc(func(ctx context.Context, m *gen.TreeNodeMutation) (ent.Value, error) {
					err := checkTxids(m)
					if err != nil {
						return nil, err
					}

					nodeIdUuid, _ := m.ID()
					nodeId := nodeIdUuid.String()

					rawTxBytes, _ := m.RawTx()
					if rawTxBytes != nil {
						rawTx, err := common.TxFromRawTxBytes(rawTxBytes)
						if err != nil {
							return nil, fmt.Errorf("failed to parse raw_tx for node %s: %w", nodeId, err)
						}
						rawTxid := rawTx.TxHash()
						m.SetRawTxid(rawTxid[:])
					}

					// Handling an update on a field
					// Check that the value was cleared. Ent framework treats this operation differently and does not show that the field is in the mutation (while in fact it is included in the final SQL), so we can not use a field getter to know if the field was cleared.
					if m.DirectTxCleared() {
						// Clear the field if it was cleared
						m.ClearDirectTxid()
					} else {
						if directTxBytes, ok := m.DirectTx(); ok {
							// If the value is set to nil directly, then it means the same as clearing the field
							if directTxBytes == nil {
								m.ClearDirectTxid()
							} else {
								directTx, err := common.TxFromRawTxBytes(directTxBytes)
								if err != nil {
									return nil, fmt.Errorf("failed to parse direct_tx for node %s: %w", nodeId, err)
								}
								directTxid := directTx.TxHash()
								m.SetDirectTxid(directTxid[:])
							}
						}
					}

					if m.DirectFromCpfpRefundTxCleared() {
						m.ClearDirectFromCpfpRefundTxid()
					} else {
						if directFromCpfpRefundTxBytes, ok := m.DirectFromCpfpRefundTx(); ok {
							if directFromCpfpRefundTxBytes == nil {
								m.ClearDirectFromCpfpRefundTxid()
							} else {
								directFromCpfpRefundTx, err := common.TxFromRawTxBytes(directFromCpfpRefundTxBytes)
								if err != nil {
									return nil, fmt.Errorf("failed to parse direct_from_cpfp_refundtx for node %s: %w", nodeId, err)
								}
								directFromCpfpRefundTxid := directFromCpfpRefundTx.TxHash()
								m.SetDirectFromCpfpRefundTxid(directFromCpfpRefundTxid[:])
							}
						}
					}

					if m.RawRefundTxCleared() {
						m.ClearRawRefundTxid()
					} else {
						if rawRefundTxBytes, ok := m.RawRefundTx(); ok {
							if rawRefundTxBytes == nil {
								m.ClearRawRefundTxid()
							} else {
								rawRefundTx, err := common.TxFromRawTxBytes(rawRefundTxBytes)
								if err != nil {
									return nil, fmt.Errorf("failed to parse raw_refund_tx for node %s: %w", nodeId, err)
								}
								rawRefundTxid := rawRefundTx.TxHash()
								m.SetRawRefundTxid(rawRefundTxid[:])
							}
						}
					}

					if m.DirectRefundTxCleared() {
						m.ClearDirectRefundTxid()
					} else {
						if directRefundTxBytes, ok := m.DirectRefundTx(); ok {
							if directRefundTxBytes == nil {
								m.ClearDirectRefundTxid()
							} else {
								directRefundTx, err := common.TxFromRawTxBytes(directRefundTxBytes)
								if err != nil {
									return nil, fmt.Errorf("failed to parse direct_refund_tx for node %s: %w", nodeId, err)
								}
								directRefundTxid := directRefundTx.TxHash()
								m.SetDirectRefundTxid(directRefundTxid[:])
							}
						}
					}
					return next.Mutate(ctx, m)
				})
			},
			ent.OpUpdateOne|ent.OpUpdate,
		),
		// Create ent hook
		hook.On(
			func(next ent.Mutator) ent.Mutator {
				return hook.TreeNodeFunc(func(ctx context.Context, m *gen.TreeNodeMutation) (ent.Value, error) {
					err := checkTxids(m)
					if err != nil {
						return nil, err
					}

					nodeIdUuid, _ := m.ID()
					nodeId := nodeIdUuid.String()

					rawTxBytes, _ := m.RawTx()
					rawTx, err := common.TxFromRawTxBytes(rawTxBytes)
					if err != nil {
						return nil, fmt.Errorf("failed to parse raw_tx for node %s: %w", nodeId, err)
					}
					rawTxid := rawTx.TxHash()
					m.SetRawTxid(rawTxid[:])

					// Handling creation is easier because we don't need to check if the value was changed.
					directTxBytes, _ := m.DirectTx()
					if directTxBytes != nil {
						directTx, err := common.TxFromRawTxBytes(directTxBytes)
						if err != nil {
							return nil, fmt.Errorf("failed to parse direct_tx for node %s: %w", nodeId, err)
						}
						directTxid := directTx.TxHash()
						m.SetDirectTxid(directTxid[:])
					}

					directFromCpfpRefundTxBytes, _ := m.DirectFromCpfpRefundTx()
					if directFromCpfpRefundTxBytes != nil {
						directFromCpfpRefundTx, err := common.TxFromRawTxBytes(directFromCpfpRefundTxBytes)
						if err != nil {
							return nil, fmt.Errorf("failed to parse direct_from_cpfp_refundtx for node %s: %w", nodeId, err)
						}
						directFromCpfpRefundTxid := directFromCpfpRefundTx.TxHash()
						m.SetDirectFromCpfpRefundTxid(directFromCpfpRefundTxid[:])
					}

					rawRefundTxBytes, _ := m.RawRefundTx()
					if rawRefundTxBytes != nil {
						rawRefundTx, err := common.TxFromRawTxBytes(rawRefundTxBytes)
						if err != nil {
							return nil, fmt.Errorf("failed to parse raw_refund_tx for node %s: %w", nodeId, err)
						}
						rawRefundTxid := rawRefundTx.TxHash()
						m.SetRawRefundTxid(rawRefundTxid[:])
					}

					directRefundTxBytes, _ := m.DirectRefundTx()
					if directRefundTxBytes != nil {
						directRefundTx, err := common.TxFromRawTxBytes(directRefundTxBytes)
						if err != nil {
							return nil, fmt.Errorf("failed to parse direct_refund_tx for node %s: %w", nodeId, err)
						}
						directRefundTxid := directRefundTx.TxHash()
						m.SetDirectRefundTxid(directRefundTxid[:])
					}
					return next.Mutate(ctx, m)
				})
			},
			ent.OpCreate,
		),
	}
}

// Checks if the txids are set directly.
// This is not allowed because the txids are generated by hooks from the transactions.
func checkTxids(m *gen.TreeNodeMutation) error {
	if _, ok := m.RawTxid(); ok {
		return fmt.Errorf("raw_txid is not allowed to be set directly")
	}
	if _, ok := m.DirectTxid(); ok {
		return fmt.Errorf("direct_txid is not allowed to be set directly")
	}
	if _, ok := m.DirectFromCpfpRefundTxid(); ok {
		return fmt.Errorf("direct_from_cpfp_refund_txid is not allowed to be set directly")
	}
	if _, ok := m.RawRefundTxid(); ok {
		return fmt.Errorf("raw_refund_txid is not allowed to be set directly")
	}
	if _, ok := m.DirectRefundTxid(); ok {
		return fmt.Errorf("direct_refund_txid is not allowed to be set directly")
	}
	return nil
}
