package vault

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"math/big"
	"strings"

	"github.com/decred/dcrd/dcrec/secp256k1/v4"
	"github.com/decred/dcrd/dcrec/secp256k1/v4/ecdsa"
	"github.com/ethereum/go-ethereum/core/types"
	troncore "github.com/fbsobreira/gotron-sdk/pkg/proto/core"
	"github.com/tech/sendico/pkg/vault/managedkey"
	"go.uber.org/zap"
	"google.golang.org/protobuf/proto"

	"github.com/tech/sendico/gateway/tron/internal/keymanager"
	"github.com/tech/sendico/pkg/merrors"
	"github.com/tech/sendico/pkg/mlogger"
)

// Config describes how to connect to Vault for managed wallet keys.
type Config = managedkey.Config

// Manager implements the keymanager.Manager contract backed by HashiCorp Vault.
type Manager struct {
	logger mlogger.Logger
	keys   managedkey.Service
}

// New constructs a Vault-backed key manager.
func New(logger mlogger.Logger, cfg Config) (*Manager, error) {
	if logger == nil {
		return nil, merrors.InvalidArgument("vault key manager: logger is required")
	}
	keys, err := managedkey.New(managedkey.Options{
		Logger:           logger,
		Config:           managedkey.Config(cfg),
		Component:        "vault key manager",
		DefaultKeyPrefix: "gateway/tron/wallets",
	})
	if err != nil {
		logger.Error("Failed to initialise vault key manager", zap.Error(err))
		return nil, err
	}

	return &Manager{
		logger: logger.Named("vault"),
		keys:   keys,
	}, nil
}

// CreateManagedWalletKey creates a new managed wallet key and stores it in Vault.
func (m *Manager) CreateManagedWalletKey(ctx context.Context, walletRef string, network string) (*keymanager.ManagedWalletKey, error) {
	created, err := m.keys.CreateManagedWalletKey(ctx, walletRef, network)
	if err != nil {
		m.logger.Warn("Failed to create managed wallet key", zap.String("wallet_ref", walletRef), zap.String("network", network), zap.Error(err))
		return nil, err
	}

	return &keymanager.ManagedWalletKey{
		KeyID:     created.KeyID,
		Address:   created.Address,
		PublicKey: created.PublicKey,
	}, nil
}

// SignTransaction loads the key material from Vault and signs the transaction.
func (m *Manager) SignTransaction(ctx context.Context, keyID string, tx *types.Transaction, chainID *big.Int) (*types.Transaction, error) {
	signed, err := m.keys.SignEVMTransaction(ctx, keyID, tx, chainID)
	if err != nil {
		m.logger.Warn("Failed to sign transaction", zap.String("key_id", keyID), zap.Error(err))
		return nil, err
	}
	return signed, nil
}

// SignTronTransaction signs a native TRON transaction using the identified key material.
func (m *Manager) SignTronTransaction(ctx context.Context, keyID string, tx *troncore.Transaction) (*troncore.Transaction, error) {
	if strings.TrimSpace(keyID) == "" {
		m.logger.Warn("TRON signing failed: empty key id")
		return nil, merrors.InvalidArgument("vault key manager: keyID is required")
	}
	if tx == nil {
		m.logger.Warn("TRON signing failed: nil transaction", zap.String("key_id", keyID))
		return nil, merrors.InvalidArgument("vault key manager: transaction is nil")
	}
	if tx.GetRawData() == nil {
		m.logger.Warn("TRON signing failed: nil raw data", zap.String("key_id", keyID))
		return nil, merrors.InvalidArgument("vault key manager: transaction raw_data is nil")
	}

	material, err := m.keys.LoadKeyMaterial(ctx, keyID)
	if err != nil {
		m.logger.Warn("Failed to load key material for TRON signing", zap.String("key_id", keyID), zap.Error(err))
		return nil, err
	}

	keyBytes, err := hex.DecodeString(material.PrivateKey)
	if err != nil {
		m.logger.Warn("Invalid key material for TRON signing", zap.String("key_id", keyID), zap.Error(err))
		return nil, merrors.Internal("vault key manager: invalid key material: " + err.Error())
	}
	defer zeroBytes(keyBytes)

	// Marshal the raw_data to bytes for hashing
	rawBytes, err := proto.Marshal(tx.GetRawData())
	if err != nil {
		m.logger.Warn("Failed to marshal TRON transaction raw_data", zap.String("key_id", keyID), zap.Error(err))
		return nil, merrors.Internal("vault key manager: failed to marshal transaction: " + err.Error())
	}

	// SHA256 hash of the raw_data
	hash := sha256.Sum256(rawBytes)

	// Create secp256k1 private key
	privKey := secp256k1.PrivKeyFromBytes(keyBytes)

	// Sign using compact signature format (65 bytes: r[32] || s[32] || recovery_id[1])
	signature := ecdsa.SignCompact(privKey, hash[:], false)

	// TRON expects signature in format: r[32] || s[32] || v[1]
	// SignCompact returns: recovery_id[1] || r[32] || s[32]
	// We need to rearrange to: r[32] || s[32] || recovery_id[1]
	if len(signature) != 65 {
		m.logger.Warn("Unexpected signature length", zap.String("key_id", keyID), zap.Int("length", len(signature)))
		return nil, merrors.Internal("vault key manager: unexpected signature length")
	}

	tronSig := make([]byte, 65)
	copy(tronSig[0:32], signature[1:33])   // r
	copy(tronSig[32:64], signature[33:65]) // s
	tronSig[64] = signature[0]             // recovery id (v)

	// Append signature to transaction
	tx.Signature = append(tx.Signature, tronSig)

	m.logger.Info("TRON transaction signed with managed key",
		zap.String("key_id", keyID),
		zap.String("network", material.Network),
	)

	return tx, nil
}

func zeroBytes(data []byte) {
	for i := range data {
		data[i] = 0
	}
}

var _ keymanager.Manager = (*Manager)(nil)