package license

import (
	"context"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"slices"
	"strings"
	"sync"
	"time"
)

const (
	credentialFormatVersion = 1
	credentialKeyBytes      = 32
	maxCredentialBytes      = 1024 * 1024
	defaultSigningKeyID     = "license-default"
)

var (
	ErrCredentialCorrupt  = errors.New("encrypted license credential is corrupt")
	ErrCredentialKey      = errors.New("license credential key is unavailable")
	ErrCredentialMismatch = errors.New("license credential instance binding mismatch")
)

type KeyMaterial struct {
	ID  string
	Key []byte
}

type KeyProvider interface {
	LoadOrCreate(context.Context) (KeyMaterial, error)
}

type FileKeyProvider struct {
	path string
	mu   sync.Mutex
}

func NewFileKeyProvider(path string) *FileKeyProvider {
	return &FileKeyProvider{path: strings.TrimSpace(path)}
}

func (p *FileKeyProvider) LoadOrCreate(ctx context.Context) (KeyMaterial, error) {
	if p == nil || p.path == "" {
		return KeyMaterial{}, fmt.Errorf("%w: key file path is empty", ErrCredentialKey)
	}
	if err := ctx.Err(); err != nil {
		return KeyMaterial{}, err
	}
	p.mu.Lock()
	defer p.mu.Unlock()

	key, err := readKeyFile(p.path)
	if err == nil {
		return keyMaterial(key), nil
	}
	if !errors.Is(err, os.ErrNotExist) {
		for range 10 {
			if err := ctx.Err(); err != nil {
				return KeyMaterial{}, err
			}
			time.Sleep(25 * time.Millisecond)
			if existing, readErr := readKeyFile(p.path); readErr == nil {
				return keyMaterial(existing), nil
			}
		}
		return KeyMaterial{}, fmt.Errorf("%w: %v", ErrCredentialKey, err)
	}

	if err := os.MkdirAll(filepath.Dir(p.path), 0o700); err != nil {
		return KeyMaterial{}, fmt.Errorf("%w: create key directory: %v", ErrCredentialKey, err)
	}
	key = make([]byte, credentialKeyBytes)
	if _, err := io.ReadFull(rand.Reader, key); err != nil {
		return KeyMaterial{}, fmt.Errorf("%w: generate key: %v", ErrCredentialKey, err)
	}

	file, err := os.OpenFile(p.path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600)
	if errors.Is(err, os.ErrExist) {
		for range 10 {
			if err := ctx.Err(); err != nil {
				return KeyMaterial{}, err
			}
			if existing, readErr := readKeyFile(p.path); readErr == nil {
				return keyMaterial(existing), nil
			}
			time.Sleep(25 * time.Millisecond)
		}
		return KeyMaterial{}, fmt.Errorf("%w: concurrent key creation did not complete", ErrCredentialKey)
	}
	if err != nil {
		return KeyMaterial{}, fmt.Errorf("%w: create key file: %v", ErrCredentialKey, err)
	}
	writeErr := func() error {
		if _, err := file.Write(key); err != nil {
			return err
		}
		return file.Sync()
	}()
	closeErr := file.Close()
	if writeErr != nil {
		return KeyMaterial{}, fmt.Errorf("%w: write key file: %v", ErrCredentialKey, writeErr)
	}
	if closeErr != nil {
		return KeyMaterial{}, fmt.Errorf("%w: close key file: %v", ErrCredentialKey, closeErr)
	}
	_ = os.Chmod(p.path, 0o600)
	return keyMaterial(key), nil
}

func readKeyFile(path string) ([]byte, error) {
	key, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	if len(key) != credentialKeyBytes {
		return nil, fmt.Errorf("invalid key length %d", len(key))
	}
	return key, nil
}

func keyMaterial(key []byte) KeyMaterial {
	digest := sha256.Sum256(key)
	return KeyMaterial{
		ID:  "file-" + hex.EncodeToString(digest[:8]),
		Key: append([]byte(nil), key...),
	}
}

type credentialEnvelope struct {
	FormatVersion          int       `json:"formatVersion"`
	Code                   string    `json:"licenseCode"`
	Type                   string    `json:"licenseType"`
	Version                string    `json:"licenseVersion"`
	VerifiedAt             time.Time `json:"verifiedAt"`
	Identity               Identity  `json:"identity"`
	InstanceToken          string    `json:"instanceToken"`
	TokenExpiresAt         time.Time `json:"tokenExpiresAt"`
	LastHeartbeatAt        time.Time `json:"lastHeartbeatAt"`
	LeaseRaw               []byte    `json:"leaseRaw"`
	LeaseSignature         string    `json:"leaseSignature"`
	LeaseScope             string    `json:"leaseScope"`
	LeaseSequence          uint64    `json:"leaseSequence"`
	Features               []string  `json:"features"`
	HighWaterMark          time.Time `json:"highWaterMark"`
	SelectiveBlockDetected bool      `json:"selectiveBlockDetected"`
}

type credentialRecord struct {
	ID              int64     `gorm:"column:id;primaryKey"`
	InstanceID      string    `gorm:"column:instance_id"`
	FormatVersion   int       `gorm:"column:format_version"`
	SigningKeyID    string    `gorm:"column:signing_key_id"`
	EncryptionKeyID string    `gorm:"column:encryption_key_id"`
	LeaseScope      string    `gorm:"column:lease_scope"`
	LeaseSequence   uint64    `gorm:"column:lease_sequence"`
	Nonce           []byte    `gorm:"column:nonce"`
	Ciphertext      []byte    `gorm:"column:ciphertext"`
	RowVersion      int64     `gorm:"column:row_version"`
	CreatedAt       time.Time `gorm:"column:created_at"`
	UpdatedAt       time.Time `gorm:"column:updated_at"`
}

func (credentialRecord) TableName() string {
	return "license_credentials"
}

func sealCredential(state State, key KeyMaterial) (credentialRecord, error) {
	if len(key.Key) != credentialKeyBytes || strings.TrimSpace(key.ID) == "" {
		return credentialRecord{}, fmt.Errorf("%w: invalid AES-256 key material", ErrCredentialKey)
	}
	if !state.Configured() || len(state.LeaseRaw) == 0 || strings.TrimSpace(state.LeaseSignature) == "" {
		return credentialRecord{}, fmt.Errorf("license credential state is incomplete")
	}
	payload, _, err := extractSignedPayloadBytes(state.LeaseRaw, state.LeaseSignature)
	if err != nil {
		return credentialRecord{}, err
	}
	lease, err := decodeLeasePayload(payload)
	if err != nil {
		return credentialRecord{}, err
	}
	if err := validateLease(lease, time.Now(), state.Identity.InstanceID, false); err != nil {
		return credentialRecord{}, err
	}
	if state.LeaseSequence != 0 && state.LeaseSequence != lease.Sequence {
		return credentialRecord{}, ErrCredentialMismatch
	}
	if state.LeaseScope != "" && state.LeaseScope != lease.Scope() {
		return credentialRecord{}, ErrCredentialMismatch
	}
	envelope := credentialEnvelope{
		FormatVersion:          credentialFormatVersion,
		Code:                   strings.TrimSpace(state.Code),
		Type:                   strings.TrimSpace(state.Type),
		Version:                strings.TrimSpace(state.Version),
		VerifiedAt:             state.VerifiedAt.UTC(),
		Identity:               state.Identity,
		InstanceToken:          strings.TrimSpace(state.InstanceToken),
		TokenExpiresAt:         state.TokenExpiresAt.UTC(),
		LastHeartbeatAt:        state.LastHeartbeatAt.UTC(),
		LeaseRaw:               append([]byte(nil), state.LeaseRaw...),
		LeaseSignature:         strings.TrimSpace(state.LeaseSignature),
		LeaseScope:             lease.Scope(),
		LeaseSequence:          lease.Sequence,
		Features:               append([]string(nil), lease.Features...),
		HighWaterMark:          state.HighWaterMark.UTC(),
		SelectiveBlockDetected: state.SelectiveBlockDetected,
	}
	plaintext, err := json.Marshal(envelope)
	if err != nil {
		return credentialRecord{}, fmt.Errorf("encode license credential: %w", err)
	}
	if len(plaintext) > maxCredentialBytes {
		return credentialRecord{}, fmt.Errorf("license credential exceeds size limit")
	}

	block, err := aes.NewCipher(key.Key)
	if err != nil {
		return credentialRecord{}, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return credentialRecord{}, err
	}
	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return credentialRecord{}, fmt.Errorf("generate license credential nonce: %w", err)
	}
	signingKeyID := signingKeyID(state.LeaseRaw)
	record := credentialRecord{
		InstanceID:      strings.TrimSpace(state.Identity.InstanceID),
		FormatVersion:   credentialFormatVersion,
		SigningKeyID:    signingKeyID,
		EncryptionKeyID: key.ID,
		LeaseScope:      lease.Scope(),
		LeaseSequence:   lease.Sequence,
		Nonce:           nonce,
		RowVersion:      1,
	}
	record.Ciphertext = gcm.Seal(nil, nonce, plaintext, credentialAAD(record))
	return record, nil
}

func openCredential(record credentialRecord, key KeyMaterial) (State, error) {
	if record.FormatVersion != credentialFormatVersion || len(key.Key) != credentialKeyBytes || record.EncryptionKeyID != key.ID || record.LeaseSequence == 0 {
		return State{}, ErrCredentialCorrupt
	}
	block, err := aes.NewCipher(key.Key)
	if err != nil {
		return State{}, ErrCredentialCorrupt
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil || len(record.Nonce) != gcm.NonceSize() || len(record.Ciphertext) == 0 {
		return State{}, ErrCredentialCorrupt
	}
	plaintext, err := gcm.Open(nil, record.Nonce, record.Ciphertext, credentialAAD(record))
	if err != nil || len(plaintext) > maxCredentialBytes {
		return State{}, ErrCredentialCorrupt
	}
	var envelope credentialEnvelope
	if err := json.Unmarshal(plaintext, &envelope); err != nil || envelope.FormatVersion != credentialFormatVersion {
		return State{}, ErrCredentialCorrupt
	}
	if strings.TrimSpace(envelope.Identity.InstanceID) != strings.TrimSpace(record.InstanceID) {
		return State{}, ErrCredentialMismatch
	}
	if signingKeyID(envelope.LeaseRaw) != record.SigningKeyID || strings.TrimSpace(envelope.LeaseSignature) == "" {
		return State{}, ErrCredentialCorrupt
	}
	if strings.TrimSpace(envelope.LeaseScope) != record.LeaseScope || envelope.LeaseSequence != record.LeaseSequence {
		return State{}, ErrCredentialCorrupt
	}
	payload, _, err := extractSignedPayloadBytes(envelope.LeaseRaw, envelope.LeaseSignature)
	if err != nil {
		return State{}, ErrCredentialCorrupt
	}
	lease, err := decodeLeasePayload(payload)
	if err != nil || validateLease(lease, time.Now(), record.InstanceID, false) != nil {
		return State{}, ErrCredentialCorrupt
	}
	if lease.Scope() != record.LeaseScope || lease.Sequence != record.LeaseSequence || !slices.Equal(lease.Features, envelope.Features) {
		return State{}, ErrCredentialCorrupt
	}
	state := State{
		Code:                   strings.TrimSpace(envelope.Code),
		Type:                   strings.TrimSpace(envelope.Type),
		Version:                strings.TrimSpace(envelope.Version),
		VerifiedAt:             envelope.VerifiedAt.UTC(),
		Identity:               envelope.Identity,
		InstanceToken:          strings.TrimSpace(envelope.InstanceToken),
		TokenExpiresAt:         envelope.TokenExpiresAt.UTC(),
		LastHeartbeatAt:        envelope.LastHeartbeatAt.UTC(),
		LeaseRaw:               append([]byte(nil), envelope.LeaseRaw...),
		LeaseSignature:         strings.TrimSpace(envelope.LeaseSignature),
		LeaseScope:             strings.TrimSpace(envelope.LeaseScope),
		LeaseSequence:          envelope.LeaseSequence,
		Features:               append([]string(nil), envelope.Features...),
		HighWaterMark:          envelope.HighWaterMark.UTC(),
		SelectiveBlockDetected: envelope.SelectiveBlockDetected,
	}
	if !state.Configured() || len(state.LeaseRaw) == 0 {
		return State{}, ErrCredentialCorrupt
	}
	return state, nil
}

func credentialAAD(record credentialRecord) []byte {
	return []byte(fmt.Sprintf(
		"qycw-license|v%d|%s|%s|%s|%s|%d",
		record.FormatVersion,
		strings.TrimSpace(record.InstanceID),
		strings.TrimSpace(record.SigningKeyID),
		strings.TrimSpace(record.EncryptionKeyID),
		strings.TrimSpace(record.LeaseScope),
		record.LeaseSequence,
	))
}

func signingKeyID(raw []byte) string {
	var payload struct {
		KeyID string `json:"keyId"`
	}
	if json.Unmarshal(raw, &payload) == nil && strings.TrimSpace(payload.KeyID) != "" {
		return strings.TrimSpace(payload.KeyID)
	}
	return defaultSigningKeyID
}
