package crypto

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"time"

	"codeberg.org/mapleopentech/monorepo/cloud/maplefile-backend/pkg/security/crypto"
)

// tryDecodeBase64 attempts to decode a base64 string using multiple encodings.
// It tries URL-safe without padding first (libsodium's URLSAFE_NO_PADDING),
// then standard base64 with padding, then standard without padding.
func tryDecodeBase64(s string) ([]byte, error) {
	var lastErr error

	// Try URL-safe base64 without padding (libsodium's URLSAFE_NO_PADDING)
	if data, err := base64.RawURLEncoding.DecodeString(s); err == nil {
		return data, nil
	} else {
		lastErr = err
	}

	// Try standard base64 with padding (Go's default for []byte)
	if data, err := base64.StdEncoding.DecodeString(s); err == nil {
		return data, nil
	} else {
		lastErr = err
	}

	// Try standard base64 without padding
	if data, err := base64.RawStdEncoding.DecodeString(s); err == nil {
		return data, nil
	} else {
		lastErr = err
	}

	// Try URL-safe base64 with padding
	if data, err := base64.URLEncoding.DecodeString(s); err == nil {
		return data, nil
	} else {
		lastErr = err
	}

	return nil, fmt.Errorf("failed to decode base64 with any encoding: %w", lastErr)
}

// MasterKey represents the root encryption key for a user
type MasterKey struct {
	Key []byte `json:"key" bson:"key"`
}

// EncryptedMasterKey is the master key encrypted with the key encryption key
type EncryptedMasterKey struct {
	Ciphertext   []byte                   `json:"ciphertext" bson:"ciphertext"`
	Nonce        []byte                   `json:"nonce" bson:"nonce"`
	KeyVersion   int                      `json:"key_version" bson:"key_version"`
	RotatedAt    *time.Time               `json:"rotated_at,omitempty" bson:"rotated_at,omitempty"`
	PreviousKeys []EncryptedHistoricalKey `json:"previous_keys,omitempty" bson:"previous_keys,omitempty"`
}

func (emk *EncryptedMasterKey) GetCurrentVersion() int {
	return emk.KeyVersion
}

func (emk *EncryptedMasterKey) GetKeyByVersion(version int) *EncryptedHistoricalKey {
	if version == emk.KeyVersion {
		// Return current key as historical format
		return &EncryptedHistoricalKey{
			KeyVersion: emk.KeyVersion,
			Ciphertext: emk.Ciphertext,
			Nonce:      emk.Nonce,
			Algorithm:  crypto.ChaCha20Poly1305Algorithm, // ✅ Updated to ChaCha20-Poly1305
		}
	}

	for _, key := range emk.PreviousKeys {
		if key.KeyVersion == version {
			return &key
		}
	}
	return nil
}

// KeyEncryptionKey derived from user password
type KeyEncryptionKey struct {
	Key  []byte `json:"key" bson:"key"`
	Salt []byte `json:"salt" bson:"salt"`
}

// PublicKey for asymmetric encryption
type PublicKey struct {
	Key            []byte `json:"key" bson:"key"`
	VerificationID string `json:"verification_id" bson:"verification_id"`
}

// PrivateKey for asymmetric decryption
type PrivateKey struct {
	Key []byte `json:"key" bson:"key"`
}

// EncryptedPrivateKey is the private key encrypted with the master key
type EncryptedPrivateKey struct {
	Ciphertext []byte `json:"ciphertext" bson:"ciphertext"`
	Nonce      []byte `json:"nonce" bson:"nonce"`
}

// RecoveryKey for account recovery
type RecoveryKey struct {
	Key []byte `json:"key" bson:"key"`
}

// EncryptedRecoveryKey is the recovery key encrypted with the master key
type EncryptedRecoveryKey struct {
	Ciphertext []byte `json:"ciphertext" bson:"ciphertext"`
	Nonce      []byte `json:"nonce" bson:"nonce"`
}

// CollectionKey encrypts files in a collection
type CollectionKey struct {
	Key          []byte `json:"key" bson:"key"`
	CollectionID string `json:"collection_id" bson:"collection_id"`
}

// EncryptedCollectionKey is the collection key encrypted with master key
type EncryptedCollectionKey struct {
	Ciphertext   []byte                   `json:"ciphertext" bson:"ciphertext"`
	Nonce        []byte                   `json:"nonce" bson:"nonce"`
	KeyVersion   int                      `json:"key_version" bson:"key_version"`
	RotatedAt    *time.Time               `json:"rotated_at,omitempty" bson:"rotated_at,omitempty"`
	PreviousKeys []EncryptedHistoricalKey `json:"previous_keys,omitempty" bson:"previous_keys,omitempty"`
}

func (eck *EncryptedCollectionKey) NeedsRotation(policy KeyRotationPolicy) bool {
	if eck.RotatedAt == nil {
		return true // Never rotated
	}

	keyAge := time.Since(*eck.RotatedAt)
	return keyAge > policy.MaxKeyAge
}

// MarshalJSON custom marshaller for EncryptedCollectionKey to serialize bytes as base64 strings.
func (eck *EncryptedCollectionKey) MarshalJSON() ([]byte, error) {
	type Alias struct {
		Ciphertext string `json:"ciphertext"`
		Nonce      string `json:"nonce"`
		KeyVersion int    `json:"key_version"`
	}
	alias := Alias{
		Ciphertext: base64.StdEncoding.EncodeToString(eck.Ciphertext),
		Nonce:      base64.StdEncoding.EncodeToString(eck.Nonce),
		KeyVersion: eck.KeyVersion,
	}
	return json.Marshal(alias)
}

// UnmarshalJSON custom unmarshaller for EncryptedCollectionKey to handle URL-safe base64 strings.
func (eck *EncryptedCollectionKey) UnmarshalJSON(data []byte) error {
	// Temporary struct to unmarshal into string fields
	type Alias struct {
		Ciphertext string `json:"ciphertext"`
		Nonce      string `json:"nonce"`
		KeyVersion int    `json:"key_version"`
	}
	var alias Alias

	if err := json.Unmarshal(data, &alias); err != nil {
		return fmt.Errorf("failed to unmarshal EncryptedCollectionKey into alias: %w", err)
	}

	// Set KeyVersion
	eck.KeyVersion = alias.KeyVersion

	// Decode Ciphertext - try multiple base64 encodings
	if alias.Ciphertext != "" {
		ciphertextBytes, err := tryDecodeBase64(alias.Ciphertext)
		if err != nil {
			return fmt.Errorf("failed to decode EncryptedCollectionKey.Ciphertext: %w", err)
		}
		eck.Ciphertext = ciphertextBytes
	}

	// Decode Nonce - try multiple base64 encodings
	if alias.Nonce != "" {
		nonceBytes, err := tryDecodeBase64(alias.Nonce)
		if err != nil {
			return fmt.Errorf("failed to decode EncryptedCollectionKey.Nonce: %w", err)
		}
		eck.Nonce = nonceBytes
	}

	return nil
}

// FileKey encrypts a specific file
type FileKey struct {
	Key    []byte `json:"key" bson:"key"`
	FileID string `json:"file_id" bson:"file_id"`
}

// EncryptedFileKey is the file key encrypted with collection key
type EncryptedFileKey struct {
	Ciphertext   []byte                   `json:"ciphertext" bson:"ciphertext"`
	Nonce        []byte                   `json:"nonce" bson:"nonce"`
	KeyVersion   int                      `json:"key_version" bson:"key_version"`
	RotatedAt    *time.Time               `json:"rotated_at,omitempty" bson:"rotated_at,omitempty"`
	PreviousKeys []EncryptedHistoricalKey `json:"previous_keys,omitempty" bson:"previous_keys,omitempty"`
}

func (eck *EncryptedFileKey) NeedsRotation(policy KeyRotationPolicy) bool {
	if eck.RotatedAt == nil {
		return true // Never rotated
	}

	keyAge := time.Since(*eck.RotatedAt)
	return keyAge > policy.MaxKeyAge
}

// MarshalJSON custom marshaller for EncryptedFileKey to serialize bytes as base64 strings.
func (efk *EncryptedFileKey) MarshalJSON() ([]byte, error) {
	type Alias struct {
		Ciphertext string `json:"ciphertext"`
		Nonce      string `json:"nonce"`
		KeyVersion int    `json:"key_version"`
	}
	alias := Alias{
		Ciphertext: base64.StdEncoding.EncodeToString(efk.Ciphertext),
		Nonce:      base64.StdEncoding.EncodeToString(efk.Nonce),
		KeyVersion: efk.KeyVersion,
	}
	return json.Marshal(alias)
}

// UnmarshalJSON custom unmarshaller for EncryptedFileKey to handle URL-safe base64 strings.
func (efk *EncryptedFileKey) UnmarshalJSON(data []byte) error {
	// Temporary struct to unmarshal into string fields
	type Alias struct {
		Ciphertext string `json:"ciphertext"`
		Nonce      string `json:"nonce"`
		KeyVersion int    `json:"key_version"`
	}
	var alias Alias

	if err := json.Unmarshal(data, &alias); err != nil {
		return fmt.Errorf("failed to unmarshal EncryptedFileKey into alias: %w", err)
	}

	// Set KeyVersion
	efk.KeyVersion = alias.KeyVersion

	// Decode Ciphertext - try multiple base64 encodings
	if alias.Ciphertext != "" {
		ciphertextBytes, err := tryDecodeBase64(alias.Ciphertext)
		if err != nil {
			return fmt.Errorf("failed to decode EncryptedFileKey.Ciphertext: %w", err)
		}
		efk.Ciphertext = ciphertextBytes
	}

	// Decode Nonce - try multiple base64 encodings
	if alias.Nonce != "" {
		nonceBytes, err := tryDecodeBase64(alias.Nonce)
		if err != nil {
			return fmt.Errorf("failed to decode EncryptedFileKey.Nonce: %w", err)
		}
		efk.Nonce = nonceBytes
	}

	return nil
}

// TagKey encrypts tag data (name and color)
type TagKey struct {
	Key    []byte `json:"key" bson:"key"`
	TagID  string `json:"tag_id" bson:"tag_id"`
}

// EncryptedTagKey is the tag key encrypted with user's master key
type EncryptedTagKey struct {
	Ciphertext   []byte                   `json:"ciphertext" bson:"ciphertext"`
	Nonce        []byte                   `json:"nonce" bson:"nonce"`
	KeyVersion   int                      `json:"key_version" bson:"key_version"`
	RotatedAt    *time.Time               `json:"rotated_at,omitempty" bson:"rotated_at,omitempty"`
	PreviousKeys []EncryptedHistoricalKey `json:"previous_keys,omitempty" bson:"previous_keys,omitempty"`
}

func (etk *EncryptedTagKey) NeedsRotation(policy KeyRotationPolicy) bool {
	if etk.RotatedAt == nil {
		return true // Never rotated
	}

	keyAge := time.Since(*etk.RotatedAt)
	return keyAge > policy.MaxKeyAge
}

// MarshalJSON custom marshaller for EncryptedTagKey to serialize bytes as base64 strings.
func (etk *EncryptedTagKey) MarshalJSON() ([]byte, error) {
	type Alias struct {
		Ciphertext string `json:"ciphertext"`
		Nonce      string `json:"nonce"`
		KeyVersion int    `json:"key_version"`
	}
	alias := Alias{
		Ciphertext: base64.StdEncoding.EncodeToString(etk.Ciphertext),
		Nonce:      base64.StdEncoding.EncodeToString(etk.Nonce),
		KeyVersion: etk.KeyVersion,
	}
	return json.Marshal(alias)
}

// UnmarshalJSON custom unmarshaller for EncryptedTagKey to handle URL-safe base64 strings.
func (etk *EncryptedTagKey) UnmarshalJSON(data []byte) error {
	// Temporary struct to unmarshal into string fields
	type Alias struct {
		Ciphertext string `json:"ciphertext"`
		Nonce      string `json:"nonce"`
		KeyVersion int    `json:"key_version"`
	}
	var alias Alias

	if err := json.Unmarshal(data, &alias); err != nil {
		return fmt.Errorf("failed to unmarshal EncryptedTagKey into alias: %w", err)
	}

	// Set KeyVersion
	etk.KeyVersion = alias.KeyVersion

	// Decode Ciphertext - try multiple base64 encodings
	if alias.Ciphertext != "" {
		ciphertextBytes, err := tryDecodeBase64(alias.Ciphertext)
		if err != nil {
			return fmt.Errorf("failed to decode EncryptedTagKey.Ciphertext: %w", err)
		}
		etk.Ciphertext = ciphertextBytes
	}

	// Decode Nonce - try multiple base64 encodings
	if alias.Nonce != "" {
		nonceBytes, err := tryDecodeBase64(alias.Nonce)
		if err != nil {
			return fmt.Errorf("failed to decode EncryptedTagKey.Nonce: %w", err)
		}
		etk.Nonce = nonceBytes
	}

	return nil
}

// MasterKeyEncryptedWithRecoveryKey allows account recovery
type MasterKeyEncryptedWithRecoveryKey struct {
	Ciphertext []byte `json:"ciphertext" bson:"ciphertext"`
	Nonce      []byte `json:"nonce" bson:"nonce"`
}