// manager, _ := node.NewNodeManager()
// manager.GenerateIdentity("my-worker", node.RoleWorker)
// identity := manager.GetIdentity()
package node

import (
	"crypto/ecdh"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"os"
	"path"
	"sync"
	"time"

	"forge.lthn.ai/Snider/Borg/pkg/stmf"
	"github.com/adrg/xdg"
)

// challenge := make([]byte, ChallengeSize)
// rand.Read(challenge)
const ChallengeSize = 32

// challenge, err := node.GenerateChallenge()
// if err != nil { return err }
// payload := HandshakePayload{Challenge: challenge}
func GenerateChallenge() ([]byte, error) {
	challenge := make([]byte, ChallengeSize)
	if _, err := rand.Read(challenge); err != nil {
		return nil, &ProtocolError{Code: ErrCodeOperationFailed, Message: "failed to generate challenge: " + err.Error()}
	}
	return challenge, nil
}

// sig := node.SignChallenge(challenge, sharedSecret)
// payload := HandshakeAckPayload{ChallengeResponse: sig}
func SignChallenge(challenge []byte, sharedSecret []byte) []byte {
	messageAuthCode := hmac.New(sha256.New, sharedSecret)
	messageAuthCode.Write(challenge)
	return messageAuthCode.Sum(nil)
}

// if !node.VerifyChallenge(challenge, ackPayload.ChallengeResponse, sharedSecret) {
//     return fmt.Errorf("challenge response verification failed")
// }
func VerifyChallenge(challenge, response, sharedSecret []byte) bool {
	expected := SignChallenge(challenge, sharedSecret)
	return hmac.Equal(response, expected)
}

// nodeManager.GenerateIdentity("my-worker", RoleWorker)
// nodeManager.GenerateIdentity("controller-1", RoleController)
// nodeManager.GenerateIdentity("fleet-node", RoleDual)
type NodeRole string

const (
	// RoleController manages remote worker nodes.
	RoleController NodeRole = "controller"
	// RoleWorker receives commands and runs miners.
	RoleWorker NodeRole = "worker"
	// RoleDual operates as both controller and worker (default).
	RoleDual NodeRole = "dual"
)

// identity := manager.GetIdentity()
// log(identity.ID, identity.Name, identity.PublicKey, identity.Role)
type NodeIdentity struct {
	ID        string    `json:"id"`        // Derived from public key (first 16 bytes hex)
	Name      string    `json:"name"`      // Human-friendly name
	PublicKey string    `json:"publicKey"` // X25519 base64
	CreatedAt time.Time `json:"createdAt"`
	Role      NodeRole  `json:"role"`
}

// manager, err := node.NewNodeManager()
// if !manager.HasIdentity() { manager.GenerateIdentity("my-worker", RoleWorker) }
// identity := manager.GetIdentity()
type NodeManager struct {
	identity   *NodeIdentity
	privateKey []byte // Never serialized to JSON
	keyPair    *stmf.KeyPair
	keyPath    string // ~/.local/share/lethean-desktop/node/private.key
	configPath string // ~/.config/lethean-desktop/node.json
	mutex sync.RWMutex
}

// nodeManager, err := node.NewNodeManager()
// if err != nil { return err }
// if !nodeManager.HasIdentity() { nodeManager.GenerateIdentity("my-node", RoleDual) }
func NewNodeManager() (*NodeManager, error) {
	keyPath, err := xdg.DataFile("lethean-desktop/node/private.key")
	if err != nil {
		return nil, &ProtocolError{Code: ErrCodeOperationFailed, Message: "failed to get key path: " + err.Error()}
	}

	configPath, err := xdg.ConfigFile("lethean-desktop/node.json")
	if err != nil {
		return nil, &ProtocolError{Code: ErrCodeOperationFailed, Message: "failed to get config path: " + err.Error()}
	}

	return NewNodeManagerWithPaths(keyPath, configPath)
}

// manager, err := node.NewNodeManagerWithPaths("/tmp/test.key", "/tmp/test.json")
// if err != nil { return err }
func NewNodeManagerWithPaths(keyPath, configPath string) (*NodeManager, error) {
	manager := &NodeManager{
		keyPath:    keyPath,
		configPath: configPath,
	}

	if err := manager.loadIdentity(); err != nil {
		// Identity doesn't exist yet, that's ok
		return manager, nil
	}

	return manager, nil
}

// if nodeManager.HasIdentity() { return nodeManager.GetIdentity() }
func (manager *NodeManager) HasIdentity() bool {
	manager.mutex.RLock()
	defer manager.mutex.RUnlock()
	return manager.identity != nil
}

// identity := nodeManager.GetIdentity()
// if identity == nil { return core.E("node", "identity not initialised", nil) }
func (manager *NodeManager) GetIdentity() *NodeIdentity {
	manager.mutex.RLock()
	defer manager.mutex.RUnlock()
	if manager.identity == nil {
		return nil
	}
	identity := *manager.identity
	return &identity
}

// if err := nodeManager.GenerateIdentity("my-worker", RoleWorker); err != nil { return err }
func (manager *NodeManager) GenerateIdentity(name string, role NodeRole) error {
	manager.mutex.Lock()
	defer manager.mutex.Unlock()

	// Generate X25519 keypair using STMF
	keyPair, err := stmf.GenerateKeyPair()
	if err != nil {
		return &ProtocolError{Code: ErrCodeOperationFailed, Message: "failed to generate keypair: " + err.Error()}
	}

	pubKeyBytes := keyPair.PublicKey()
	hash := sha256.Sum256(pubKeyBytes)
	nodeID := hex.EncodeToString(hash[:16])

	manager.identity = &NodeIdentity{
		ID:        nodeID,
		Name:      name,
		PublicKey: keyPair.PublicKeyBase64(),
		CreatedAt: time.Now(),
		Role:      role,
	}

	manager.keyPair = keyPair
	manager.privateKey = keyPair.PrivateKey()

	// Save private key
	if err := manager.savePrivateKey(); err != nil {
		return &ProtocolError{Code: ErrCodeOperationFailed, Message: "failed to save private key: " + err.Error()}
	}

	// Save identity config
	if err := manager.saveIdentity(); err != nil {
		return &ProtocolError{Code: ErrCodeOperationFailed, Message: "failed to save identity: " + err.Error()}
	}

	return nil
}

// secret, err := nodeManager.DeriveSharedSecret(peer.PublicKey)
// smsg.Encrypt(msg, base64.StdEncoding.EncodeToString(secret))
func (manager *NodeManager) DeriveSharedSecret(peerPubKeyBase64 string) ([]byte, error) {
	manager.mutex.RLock()
	defer manager.mutex.RUnlock()

	if manager.privateKey == nil {
		return nil, &ProtocolError{Code: ErrCodeOperationFailed, Message: "node identity not initialized"}
	}

	// Load peer's public key
	peerPubKey, err := stmf.LoadPublicKeyBase64(peerPubKeyBase64)
	if err != nil {
		return nil, &ProtocolError{Code: ErrCodeOperationFailed, Message: "failed to load peer public key: " + err.Error()}
	}

	// Load our private key
	privateKey, err := ecdh.X25519().NewPrivateKey(manager.privateKey)
	if err != nil {
		return nil, &ProtocolError{Code: ErrCodeOperationFailed, Message: "failed to load private key: " + err.Error()}
	}

	// Derive shared secret using ECDH
	sharedSecret, err := privateKey.ECDH(peerPubKey)
	if err != nil {
		return nil, &ProtocolError{Code: ErrCodeOperationFailed, Message: "failed to derive shared secret: " + err.Error()}
	}

	hash := sha256.Sum256(sharedSecret)
	return hash[:], nil
}

// manager.savePrivateKey() // called from GenerateIdentity after keypair generation
func (manager *NodeManager) savePrivateKey() error {
	// Ensure directory exists
	directoryPath := path.Dir(manager.keyPath)
	if err := os.MkdirAll(directoryPath, 0700); err != nil {
		return &ProtocolError{Code: ErrCodeOperationFailed, Message: "failed to create key directory: " + err.Error()}
	}

	// Write private key with restricted permissions (0600)
	if err := os.WriteFile(manager.keyPath, manager.privateKey, 0600); err != nil {
		return &ProtocolError{Code: ErrCodeOperationFailed, Message: "failed to write private key: " + err.Error()}
	}

	return nil
}

// manager.saveIdentity() // called from GenerateIdentity after savePrivateKey succeeds
func (manager *NodeManager) saveIdentity() error {
	// Ensure directory exists
	directoryPath := path.Dir(manager.configPath)
	if err := os.MkdirAll(directoryPath, 0755); err != nil {
		return &ProtocolError{Code: ErrCodeOperationFailed, Message: "failed to create config directory: " + err.Error()}
	}

	data, err := MarshalJSON(manager.identity)
	if err != nil {
		return &ProtocolError{Code: ErrCodeOperationFailed, Message: "failed to marshal identity: " + err.Error()}
	}

	if err := os.WriteFile(manager.configPath, data, 0644); err != nil {
		return &ProtocolError{Code: ErrCodeOperationFailed, Message: "failed to write identity: " + err.Error()}
	}

	return nil
}

// manager.loadIdentity() // called from NewNodeManagerWithPaths on startup; nil error means identity is ready
func (manager *NodeManager) loadIdentity() error {
	// Load identity config
	data, err := os.ReadFile(manager.configPath)
	if err != nil {
		return &ProtocolError{Code: ErrCodeOperationFailed, Message: "failed to read identity: " + err.Error()}
	}

	var identity NodeIdentity
	if err := UnmarshalJSON(data, &identity); err != nil {
		return &ProtocolError{Code: ErrCodeOperationFailed, Message: "failed to unmarshal identity: " + err.Error()}
	}

	// Load private key
	privateKey, err := os.ReadFile(manager.keyPath)
	if err != nil {
		return &ProtocolError{Code: ErrCodeOperationFailed, Message: "failed to read private key: " + err.Error()}
	}

	// Reconstruct keypair from private key
	keyPair, err := stmf.LoadKeyPair(privateKey)
	if err != nil {
		return &ProtocolError{Code: ErrCodeOperationFailed, Message: "failed to load keypair: " + err.Error()}
	}

	manager.identity = &identity
	manager.privateKey = privateKey
	manager.keyPair = keyPair

	return nil
}

// if err := manager.Delete(); err != nil { return err } // removes key and config from disk; clears in-memory state
func (manager *NodeManager) Delete() error {
	manager.mutex.Lock()
	defer manager.mutex.Unlock()

	// Remove private key
	if err := os.Remove(manager.keyPath); err != nil && !os.IsNotExist(err) {
		return &ProtocolError{Code: ErrCodeOperationFailed, Message: "failed to remove private key: " + err.Error()}
	}

	// Remove identity config
	if err := os.Remove(manager.configPath); err != nil && !os.IsNotExist(err) {
		return &ProtocolError{Code: ErrCodeOperationFailed, Message: "failed to remove identity: " + err.Error()}
	}

	manager.identity = nil
	manager.privateKey = nil
	manager.keyPair = nil

	return nil
}