// Package seal implements the seal/unseal state machine for Metacrypt.
package seal

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"log/slog"
	"sync"
	"time"

	"git.wntrmute.dev/kyle/metacrypt/internal/barrier"
	"git.wntrmute.dev/kyle/metacrypt/internal/crypto"
)

// ServiceState represents the current state of the Metacrypt service.
type ServiceState int

const (
	StateUninitialized ServiceState = iota
	StateSealed
	StateInitializing
	StateUnsealed
)

func (s ServiceState) String() string {
	switch s {
	case StateUninitialized:
		return "uninitialized"
	case StateSealed:
		return "sealed"
	case StateInitializing:
		return "initializing"
	case StateUnsealed:
		return "unsealed"
	default:
		return "unknown"
	}
}

var (
	ErrAlreadyInitialized = errors.New("seal: already initialized")
	ErrNotInitialized     = errors.New("seal: not initialized")
	ErrInvalidPassword    = errors.New("seal: invalid password")
	ErrSealed             = errors.New("seal: service is sealed")
	ErrNotSealed          = errors.New("seal: service is not sealed")
	ErrRateLimited        = errors.New("seal: too many unseal attempts, try again later")
)

// Manager manages the seal/unseal lifecycle.
type Manager struct {
	lastAttempt    time.Time
	lockoutUntil   time.Time
	db             *sql.DB
	barrier        *barrier.AESGCMBarrier
	logger         *slog.Logger
	mek            []byte
	state          ServiceState
	unsealAttempts int
	mu             sync.RWMutex
}

// NewManager creates a new seal manager.
func NewManager(db *sql.DB, b *barrier.AESGCMBarrier, logger *slog.Logger) *Manager {
	return &Manager{
		db:      db,
		barrier: b,
		logger:  logger,
		state:   StateUninitialized,
	}
}

// Barrier returns the underlying barrier for direct access by subsystems
// that need to read/write encrypted storage (e.g. ACME state).
// The barrier must only be used when the service is unsealed.
func (m *Manager) Barrier() *barrier.AESGCMBarrier {
	return m.barrier
}

// State returns the current service state.
func (m *Manager) State() ServiceState {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return m.state
}

// CheckInitialized checks the database for an existing seal config and
// updates the state accordingly. Should be called on startup.
func (m *Manager) CheckInitialized() error {
	m.mu.Lock()
	defer m.mu.Unlock()

	var count int
	err := m.db.QueryRow("SELECT COUNT(*) FROM seal_config").Scan(&count)
	if err != nil {
		return fmt.Errorf("seal: check initialized: %w", err)
	}
	if count > 0 {
		m.state = StateSealed
		m.logger.Debug("seal config found, state set to sealed")
	} else {
		m.state = StateUninitialized
		m.logger.Debug("no seal config found, state set to uninitialized")
	}
	return nil
}

// Initialize performs first-time setup: generates MEK, encrypts it with the
// password-derived KWK, and stores everything in seal_config.
func (m *Manager) Initialize(ctx context.Context, password []byte, params crypto.Argon2Params) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if m.state != StateUninitialized {
		return ErrAlreadyInitialized
	}

	m.logger.Debug("initializing seal manager")
	m.state = StateInitializing
	defer func() {
		if m.mek == nil {
			// If we failed, go back to uninitialized.
			m.state = StateUninitialized
		}
	}()

	// Generate salt and MEK.
	salt, err := crypto.GenerateSalt()
	if err != nil {
		return fmt.Errorf("seal: generate salt: %w", err)
	}

	mek, err := crypto.GenerateKey()
	if err != nil {
		return fmt.Errorf("seal: generate mek: %w", err)
	}

	// Derive KWK from password.
	kwk := crypto.DeriveKey(password, salt, params)
	defer crypto.Zeroize(kwk)

	// Encrypt MEK with KWK.
	encryptedMEK, err := crypto.Encrypt(kwk, mek, nil)
	if err != nil {
		crypto.Zeroize(mek)
		return fmt.Errorf("seal: encrypt mek: %w", err)
	}

	// Store in database.
	_, err = m.db.ExecContext(ctx, `
		INSERT INTO seal_config (id, encrypted_mek, kdf_salt, argon2_time, argon2_memory, argon2_threads)
		VALUES (1, ?, ?, ?, ?, ?)`,
		encryptedMEK, salt, params.Time, params.Memory, params.Threads)
	if err != nil {
		crypto.Zeroize(mek)
		return fmt.Errorf("seal: store config: %w", err)
	}

	// Unseal the barrier with the MEK.
	if err := m.barrier.Unseal(mek); err != nil {
		crypto.Zeroize(mek)
		return fmt.Errorf("seal: unseal barrier: %w", err)
	}

	m.mek = mek
	m.state = StateUnsealed
	m.logger.Debug("seal initialization complete, barrier unsealed")
	return nil
}

// Unseal decrypts the MEK using the provided password and unseals the barrier.
func (m *Manager) Unseal(password []byte) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if m.state == StateUninitialized {
		return ErrNotInitialized
	}
	if m.state == StateUnsealed {
		return ErrNotSealed
	}

	m.logger.Debug("unseal attempt")
	// Rate limiting.
	now := time.Now()
	if now.Before(m.lockoutUntil) {
		m.logger.Debug("unseal attempt rate limited")
		return ErrRateLimited
	}
	if now.Sub(m.lastAttempt) > time.Minute {
		m.unsealAttempts = 0
	}
	m.unsealAttempts++
	m.lastAttempt = now
	if m.unsealAttempts > 5 {
		m.lockoutUntil = now.Add(60 * time.Second)
		m.unsealAttempts = 0
		m.logger.Debug("unseal attempts exceeded, locking out")
		return ErrRateLimited
	}

	// Read seal config.
	var (
		encryptedMEK    []byte
		salt            []byte
		argTime, argMem uint32
		argThreads      uint8
	)
	err := m.db.QueryRow(`
		SELECT encrypted_mek, kdf_salt, argon2_time, argon2_memory, argon2_threads
		FROM seal_config WHERE id = 1`).Scan(&encryptedMEK, &salt, &argTime, &argMem, &argThreads)
	if err != nil {
		return fmt.Errorf("seal: read config: %w", err)
	}

	params := crypto.Argon2Params{Time: argTime, Memory: argMem, Threads: argThreads}

	// Derive KWK and decrypt MEK.
	kwk := crypto.DeriveKey(password, salt, params)
	defer crypto.Zeroize(kwk)

	mek, err := crypto.Decrypt(kwk, encryptedMEK, nil)
	if err != nil {
		m.logger.Debug("unseal failed: invalid password")
		return ErrInvalidPassword
	}

	// Unseal the barrier.
	if err := m.barrier.Unseal(mek); err != nil {
		crypto.Zeroize(mek)
		return fmt.Errorf("seal: unseal barrier: %w", err)
	}

	m.mek = mek
	m.state = StateUnsealed
	m.unsealAttempts = 0
	m.logger.Debug("unseal succeeded, barrier unsealed")
	return nil
}

// RotateMEK generates a new MEK, re-wraps all DEKs in the barrier, and
// updates the encrypted MEK in seal_config. The password is required to
// derive the KWK for re-encrypting the new MEK.
func (m *Manager) RotateMEK(ctx context.Context, password []byte) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if m.state != StateUnsealed {
		return ErrSealed
	}

	// Read seal config for KDF params.
	var (
		salt            []byte
		argTime, argMem uint32
		argThreads      uint8
	)
	err := m.db.QueryRow(`
		SELECT kdf_salt, argon2_time, argon2_memory, argon2_threads
		FROM seal_config WHERE id = 1`).Scan(&salt, &argTime, &argMem, &argThreads)
	if err != nil {
		return fmt.Errorf("seal: read config: %w", err)
	}

	// Verify password by decrypting existing MEK.
	params := crypto.Argon2Params{Time: argTime, Memory: argMem, Threads: argThreads}
	kwk := crypto.DeriveKey(password, salt, params)
	defer crypto.Zeroize(kwk)

	var encryptedMEK []byte
	err = m.db.QueryRow("SELECT encrypted_mek FROM seal_config WHERE id = 1").Scan(&encryptedMEK)
	if err != nil {
		return fmt.Errorf("seal: read encrypted mek: %w", err)
	}
	_, err = crypto.Decrypt(kwk, encryptedMEK, nil)
	if err != nil {
		return ErrInvalidPassword
	}

	// Generate new MEK.
	newMEK, err := crypto.GenerateKey()
	if err != nil {
		return fmt.Errorf("seal: generate new mek: %w", err)
	}

	// Encrypt new MEK with KWK before starting the transaction.
	newEncMEK, err := crypto.Encrypt(kwk, newMEK, nil)
	if err != nil {
		crypto.Zeroize(newMEK)
		return fmt.Errorf("seal: encrypt new mek: %w", err)
	}

	// Re-wrap DEKs and update seal_config in a single atomic transaction.
	tx, err := m.db.BeginTx(ctx, nil)
	if err != nil {
		crypto.Zeroize(newMEK)
		return fmt.Errorf("seal: begin tx: %w", err)
	}
	defer func() { _ = tx.Rollback() }()

	if err := m.barrier.ReWrapKeysTx(ctx, tx, newMEK); err != nil {
		crypto.Zeroize(newMEK)
		return fmt.Errorf("seal: re-wrap keys: %w", err)
	}

	_, err = tx.ExecContext(ctx,
		"UPDATE seal_config SET encrypted_mek = ? WHERE id = 1", newEncMEK)
	if err != nil {
		crypto.Zeroize(newMEK)
		return fmt.Errorf("seal: update seal config: %w", err)
	}

	if err := tx.Commit(); err != nil {
		crypto.Zeroize(newMEK)
		return fmt.Errorf("seal: commit mek rotation: %w", err)
	}

	// Only after commit: swap in-memory state.
	m.barrier.SwapMEK(newMEK)
	crypto.Zeroize(m.mek)
	m.mek = newMEK
	m.logger.Info("MEK rotated successfully")
	return nil
}

// Seal seals the service: zeroizes MEK, seals the barrier.
func (m *Manager) Seal() error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if m.state != StateUnsealed {
		return ErrNotSealed
	}

	m.logger.Debug("sealing service")
	if m.mek != nil {
		crypto.Zeroize(m.mek)
		m.mek = nil
	}
	_ = m.barrier.Seal()
	m.state = StateSealed
	m.logger.Debug("service sealed")
	return nil
}