package contracts

import (
	"crypto/ecdsa"
	"encoding/hex"
	"fmt"
	"os"

	"github.com/ethereum/go-ethereum/crypto"
)

type KeyManager struct {
	privateKey *ecdsa.PrivateKey
}

// NewKeyManager creates a new key manager from environment
func NewKeyManager() (*KeyManager, error) {
	privateKeyStr := os.Getenv("PRIVATE_KEY")
	if privateKeyStr == "" {
		return nil, fmt.Errorf("PRIVATE_KEY environment variable not set")
	}

	// Remove 0x prefix if present
	if len(privateKeyStr) > 2 && privateKeyStr[:2] == "0x" {
		privateKeyStr = privateKeyStr[2:]
	}

	privateKeyBytes, err := hex.DecodeString(privateKeyStr)
	if err != nil {
		return nil, fmt.Errorf("invalid private key: %v", err)
	}

	privateKey, err := crypto.ToECDSA(privateKeyBytes)
	if err != nil {
		return nil, fmt.Errorf("invalid private key format: %v", err)
	}

	return &KeyManager{
		privateKey: privateKey,
	}, nil
}

// GetPrivateKey returns the private key
func (km *KeyManager) GetPrivateKey() (*ecdsa.PrivateKey, error) {
	if km.privateKey == nil {
		return nil, fmt.Errorf("private key not initialized")
	}
	return km.privateKey, nil
}

// GetAddress returns the Ethereum address
func (km *KeyManager) GetAddress() string {
	if km.privateKey == nil {
		return ""
	}
	return crypto.PubkeyToAddress(km.privateKey.PublicKey).Hex()
}