package relay

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"log"
	"time"

	pubsub "github.com/libp2p/go-libp2p-pubsub"
	"github.com/libp2p/go-libp2p/core/host"

	"go-blockchain/blockchain"
	"go-blockchain/identity"
)

const (
	// TopicRelay is the gossipsub topic for encrypted envelope routing.
	TopicRelay = "dchain/relay/v1"
)

// DeliverFunc is called when a message addressed to this node is decrypted.
type DeliverFunc func(envelopeID, senderEd25519PubKey string, msg []byte)

// SubmitTxFunc submits a signed transaction to the local node's mempool.
type SubmitTxFunc func(*blockchain.Transaction) error

// Router manages relay envelope routing for a single node.
// It subscribes to TopicRelay, decrypts messages addressed to this node,
// and submits RELAY_PROOF transactions to claim delivery fees from senders.
//
// All received envelopes (regardless of recipient) are stored in the Mailbox
// so offline recipients can pull them later via GET /relay/inbox.
type Router struct {
	h         host.Host
	topic     *pubsub.Topic
	sub       *pubsub.Subscription
	kp        *KeyPair           // X25519 keypair for envelope encryption
	id        *identity.Identity // Ed25519 identity for signing RELAY_PROOF txs
	mailbox   *Mailbox           // nil disables mailbox storage
	onDeliver DeliverFunc
	submitTx  SubmitTxFunc
}

// NewRouter creates and starts a relay Router.
// mailbox may be nil to disable offline message storage.
func NewRouter(
	h host.Host,
	ps *pubsub.PubSub,
	kp *KeyPair,
	id *identity.Identity,
	mailbox *Mailbox,
	onDeliver DeliverFunc,
	submitTx SubmitTxFunc,
) (*Router, error) {
	topic, err := ps.Join(TopicRelay)
	if err != nil {
		return nil, fmt.Errorf("join relay topic: %w", err)
	}
	sub, err := topic.Subscribe()
	if err != nil {
		topic.Close()
		return nil, fmt.Errorf("subscribe relay topic: %w", err)
	}
	return &Router{
		h:         h,
		topic:     topic,
		sub:       sub,
		kp:        kp,
		id:        id,
		mailbox:   mailbox,
		onDeliver: onDeliver,
		submitTx:  submitTx,
	}, nil
}

// Send seals msg for recipientPub and broadcasts it on the relay topic.
// recipientPubHex is the hex X25519 public key of the recipient.
// feeUT is the delivery fee offered to the relay (0 = free delivery).
// Returns the envelope ID on success.
func (r *Router) Send(recipientPubHex string, msg []byte, feeUT uint64) (string, error) {
	recipBytes, err := hex.DecodeString(recipientPubHex)
	if err != nil || len(recipBytes) != 32 {
		return "", fmt.Errorf("invalid recipient pub key: %s", recipientPubHex)
	}
	var recipPub [32]byte
	copy(recipPub[:], recipBytes)

	env, err := Seal(r.kp, r.id, recipPub, msg, feeUT, time.Now().Unix())
	if err != nil {
		return "", fmt.Errorf("seal envelope: %w", err)
	}
	data, err := json.Marshal(env)
	if err != nil {
		return "", err
	}
	if err := r.topic.Publish(context.Background(), data); err != nil {
		return "", err
	}
	return env.ID, nil
}

// Run processes incoming relay envelopes until ctx is cancelled.
// Envelopes addressed to this node are decrypted and acknowledged; others
// are ignored (gossipsub handles propagation automatically).
func (r *Router) Run(ctx context.Context) {
	for {
		m, err := r.sub.Next(ctx)
		if err != nil {
			return
		}
		if m.ReceivedFrom == r.h.ID() {
			continue
		}
		var env Envelope
		if err := json.Unmarshal(m.Data, &env); err != nil {
			continue
		}

		// Store every valid envelope in the mailbox for offline delivery.
		// Messages are encrypted — the relay cannot read them.
		// ErrEnvelopeTooLarge is silently dropped (anti-spam); other errors are logged.
		if r.mailbox != nil {
			if err := r.mailbox.Store(&env); err != nil {
				if err == ErrEnvelopeTooLarge {
					log.Printf("[relay] dropped oversized envelope %s (%d bytes ciphertext)",
						env.ID, len(env.Ciphertext))
				} else {
					log.Printf("[relay] mailbox store error for %s: %v", env.ID, err)
				}
			}
		}

		if !env.IsAddressedTo(r.kp) {
			continue
		}
		msg, err := Open(r.kp, &env)
		if err != nil {
			log.Printf("[relay] decryption error for envelope %s: %v", env.ID, err)
			continue
		}
		if r.onDeliver != nil {
			r.onDeliver(env.ID, env.SenderEd25519PubKey, msg)
		}
		if r.submitTx != nil && env.FeeUT > 0 && env.SenderEd25519PubKey != "" {
			if err := r.submitRelayProof(&env); err != nil {
				log.Printf("[relay] relay proof submission failed for %s: %v", env.ID, err)
			}
		}
	}
}

// Broadcast stores env in the mailbox and publishes it on the relay gossipsub topic.
// Used by the HTTP API so light clients can send pre-sealed envelopes without
// needing a direct libp2p connection.
func (r *Router) Broadcast(env *Envelope) error {
	if r.mailbox != nil {
		if err := r.mailbox.Store(env); err != nil && err != ErrEnvelopeTooLarge {
			log.Printf("[relay] broadcast mailbox store error %s: %v", env.ID, err)
		}
	}
	data, err := json.Marshal(env)
	if err != nil {
		return err
	}
	return r.topic.Publish(context.Background(), data)
}

// RelayPubHex returns the hex X25519 public key for this node's relay keypair.
func (r *Router) RelayPubHex() string {
	return r.kp.PubHex()
}

// submitRelayProof builds and submits a RELAY_PROOF tx to claim the delivery fee.
func (r *Router) submitRelayProof(env *Envelope) error {
	envHash := Hash(env)
	relayPubKey := r.id.PubKeyHex()

	// Recipient signs envelope hash — proves the message was actually decrypted.
	recipientSig := r.id.Sign(envHash)

	payload := blockchain.RelayProofPayload{
		EnvelopeID:   env.ID,
		EnvelopeHash: envHash,
		SenderPubKey: env.SenderEd25519PubKey,
		FeeUT:        env.FeeUT,
		FeeSig:       env.FeeSig,
		RelayPubKey:  relayPubKey,
		DeliveredAt:  time.Now().Unix(),
		RecipientSig: recipientSig,
	}
	payloadBytes, err := json.Marshal(payload)
	if err != nil {
		return err
	}

	idBytes := sha256.Sum256(append([]byte(relayPubKey), []byte(env.ID)...))
	now := time.Now().UTC()
	tx := &blockchain.Transaction{
		ID:        hex.EncodeToString(idBytes[:16]),
		Type:      blockchain.EventRelayProof,
		From:      relayPubKey,
		Fee:       blockchain.MinFee,
		Payload:   payloadBytes,
		Timestamp: now,
	}
	tx.Signature = r.id.Sign(txSignBytes(tx))
	return r.submitTx(tx)
}

// txSignBytes returns the canonical bytes that are signed for a transaction,
// matching the format used in identity.txSignBytes.
func txSignBytes(tx *blockchain.Transaction) []byte {
	data, _ := json.Marshal(struct {
		ID        string               `json:"id"`
		Type      blockchain.EventType `json:"type"`
		From      string               `json:"from"`
		To        string               `json:"to"`
		Amount    uint64               `json:"amount"`
		Fee       uint64               `json:"fee"`
		Payload   []byte               `json:"payload"`
		Timestamp time.Time            `json:"timestamp"`
	}{tx.ID, tx.Type, tx.From, tx.To, tx.Amount, tx.Fee, tx.Payload, tx.Timestamp})
	return data
}