// SPDX-License-Identifier: GPL-3.0-or-later

// Package keyfile manages backup key files for Private Archive.
// A backup key file contains everything needed to restore an account:
// - Client private key (ECDSA P-256)
// - Client certificate (signed by server CA)
// - Server CA certificate
//
// The backup key file IS the only way to recover data - there is no password reset.
package keyfile

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/sha256"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"os"
	"strings"
	"time"

	"golang.org/x/crypto/argon2"
)

// BackupKeyFile represents the contents of a backup key file.
// This single file contains everything needed to access backups.
type BackupKeyFile struct {
	// PrivateKey is the ECDSA P-256 private key
	PrivateKey *ecdsa.PrivateKey

	// Certificate is the client certificate signed by the server CA
	Certificate *x509.Certificate

	// CACert is the server's CA certificate for verification
	CACert *x509.Certificate

	// CertHash is the SHA256 hash of the certificate (client identity)
	CertHash string

	// S3AccessKeyID is the AWS Sig V4 access key for S3 data operations (e.g., "pb_<cert_hash[:20]>")
	S3AccessKeyID string

	// S3SecretAccessKey is the AWS Sig V4 secret key for S3 data operations (64 hex chars, rotatable)
	S3SecretAccessKey string

	// S3BucketName is the S3 bucket assigned to this client (e.g., "pb-<cert_hash[:20]>")
	S3BucketName string

	// CreatedAt is when the key file was created
	CreatedAt time.Time
}

// s3CredentialsJSON is the JSON structure stored in the S3 credentials PEM block.
type s3CredentialsJSON struct {
	AccessKeyID     string `json:"access_key_id"`
	SecretAccessKey string `json:"secret_access_key"`
	BucketName      string `json:"bucket_name"`
}

// HasS3Credentials returns true if the key file contains S3 credentials.
func (bkf *BackupKeyFile) HasS3Credentials() bool {
	return bkf.S3AccessKeyID != "" && bkf.S3SecretAccessKey != "" && bkf.S3BucketName != ""
}

// UpdateS3Secret updates the S3 secret access key after a credential rotation.
// The access key ID and bucket name are immutable (derived from cert hash).
func (bkf *BackupKeyFile) UpdateS3Secret(newSecret string) error {
	if newSecret == "" {
		return errors.New("new secret access key cannot be empty")
	}
	bkf.S3SecretAccessKey = newSecret
	return nil
}

// Manager handles backup key file operations.
type Manager struct {
	keyFilePath string
}

// NewManager creates a new key file manager.
func NewManager(keyFilePath string) *Manager {
	return &Manager{
		keyFilePath: keyFilePath,
	}
}

// KeyFilePath returns the path to the key file.
func (m *Manager) KeyFilePath() string {
	return m.keyFilePath
}

// Exists checks if a backup key file exists.
func (m *Manager) Exists() bool {
	_, err := os.Stat(m.keyFilePath)
	return err == nil
}

// GenerateKeyPair generates a new ECDSA P-256 key pair.
func GenerateKeyPair() (*ecdsa.PrivateKey, error) {
	return ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
}

// GenerateCSR creates a Certificate Signing Request from a private key.
// Uses a random CommonName to avoid any PII.
func GenerateCSR(privateKey *ecdsa.PrivateKey) ([]byte, error) {
	// Generate random identifier for CN (no PII)
	randomBytes := make([]byte, 8)
	if _, err := rand.Read(randomBytes); err != nil {
		return nil, fmt.Errorf("failed to generate random bytes: %w", err)
	}
	randomCN := fmt.Sprintf("client-%x", randomBytes)

	template := &x509.CertificateRequest{
		Subject: pkix.Name{
			CommonName: randomCN,
		},
		SignatureAlgorithm: x509.ECDSAWithSHA256,
	}

	csrDER, err := x509.CreateCertificateRequest(rand.Reader, template, privateKey)
	if err != nil {
		return nil, fmt.Errorf("failed to create CSR: %w", err)
	}

	// Encode as PEM
	csrPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE REQUEST",
		Bytes: csrDER,
	})

	return csrPEM, nil
}

// ComputeCertHash computes the SHA256 hash of a certificate.
// This hash is used as the client identity.
func ComputeCertHash(cert *x509.Certificate) string {
	hash := sha256.Sum256(cert.Raw)
	return hex.EncodeToString(hash[:])
}

// ParseCertificatePEM parses a PEM-encoded certificate.
func ParseCertificatePEM(pemData string) (*x509.Certificate, error) {
	block, _ := pem.Decode([]byte(pemData))
	if block == nil {
		return nil, errors.New("failed to decode PEM certificate")
	}
	if block.Type != "CERTIFICATE" {
		return nil, fmt.Errorf("unexpected PEM type: %s", block.Type)
	}
	cert, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("failed to parse certificate: %w", err)
	}
	return cert, nil
}

// DeriveRepositoryPassword derives the Restic repository password from the certificate
// and private key. This ensures the user only needs one file (the backup key file) to
// recover everything.
//
// The password is derived deterministically using Argon2id with:
// - Certificate hash (identity binding)
// - Private key scalar (secret entropy)
//
// IMPORTANT: This function MUST produce the same output for the same inputs.
// The password is used to encrypt/decrypt the restic repository.
func DeriveRepositoryPassword(cert *x509.Certificate, privateKey *ecdsa.PrivateKey) (string, error) {
	// Hash the certificate DER bytes
	certHash := sha256.Sum256(cert.Raw)

	// Use the private key's scalar (D) as deterministic entropy
	// This is deterministic - same key always gives same bytes
	privateKeyBytes := privateKey.D.Bytes()

	// Pad to 32 bytes (P-256 curve order is 32 bytes)
	paddedKey := make([]byte, 32)
	copy(paddedKey[32-len(privateKeyBytes):], privateKeyBytes)

	// Combine cert hash and private key bytes
	combined := append(certHash[:], paddedKey...)

	// Use Argon2id to derive final password
	// Parameters: time=3, memory=64MB, threads=4, keyLen=32
	passwordBytes := argon2.IDKey(combined, certHash[:16], 3, 64*1024, 4, 32)

	// Base64 encode for Restic compatibility
	return base64.StdEncoding.EncodeToString(passwordBytes), nil
}

// New creates a new BackupKeyFile from component parts.
// This is typically called after receiving the signed certificate from the server.
func New(privateKey *ecdsa.PrivateKey, cert *x509.Certificate, caCert *x509.Certificate) *BackupKeyFile {
	return &BackupKeyFile{
		PrivateKey:  privateKey,
		Certificate: cert,
		CACert:      caCert,
		CertHash:    ComputeCertHash(cert),
		CreatedAt:   time.Now(),
	}
}

// NewWithS3 creates a new BackupKeyFile with S3 credentials.
// This is used during registration when the server returns S3 credentials.
func NewWithS3(privateKey *ecdsa.PrivateKey, cert *x509.Certificate, caCert *x509.Certificate,
	s3AccessKeyID, s3SecretAccessKey, s3BucketName string) *BackupKeyFile {
	return &BackupKeyFile{
		PrivateKey:        privateKey,
		Certificate:       cert,
		CACert:            caCert,
		CertHash:          ComputeCertHash(cert),
		S3AccessKeyID:     s3AccessKeyID,
		S3SecretAccessKey: s3SecretAccessKey,
		S3BucketName:      s3BucketName,
		CreatedAt:         time.Now(),
	}
}

// Save writes the backup key file to disk.
func (bkf *BackupKeyFile) Save(path string) error {
	// Encode private key
	keyBytes, err := x509.MarshalECPrivateKey(bkf.PrivateKey)
	if err != nil {
		return fmt.Errorf("failed to marshal private key: %w", err)
	}
	keyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "EC PRIVATE KEY",
		Bytes: keyBytes,
	})

	// Encode certificate
	certPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE",
		Bytes: bkf.Certificate.Raw,
	})

	// Encode CA certificate with header to distinguish it
	caPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE",
		Bytes: bkf.CACert.Raw,
		Headers: map[string]string{
			"X-Purpose": "CA",
		},
	})

	// Encode S3 credentials if present
	var s3PEM []byte
	if bkf.HasS3Credentials() {
		s3Creds := s3CredentialsJSON{
			AccessKeyID:     bkf.S3AccessKeyID,
			SecretAccessKey: bkf.S3SecretAccessKey,
			BucketName:      bkf.S3BucketName,
		}
		s3JSON, err := json.Marshal(s3Creds)
		if err != nil {
			return fmt.Errorf("failed to marshal S3 credentials: %w", err)
		}
		s3PEM = pem.EncodeToMemory(&pem.Block{
			Type:  "PRIVATEBACKUP S3 CREDENTIALS",
			Bytes: s3JSON,
		})
	}

	// Combine with header comments
	content := fmt.Sprintf(`# Private Archive Key File
# Generated: %s
# Certificate Hash: %s
#
# ⚠️  KEEP THIS FILE SAFE!
# This is the ONLY way to access your backups.
# There is NO password reset or account recovery without this file.
# Store copies in multiple secure locations.
#
# To restore on a new device:
#   privatebackup restore-account --key-file privatebackup-key.pem

%s
%s
%s`,
		bkf.CreatedAt.Format(time.RFC3339),
		bkf.CertHash,
		string(keyPEM),
		string(certPEM),
		string(caPEM),
	)

	// Append S3 credentials block if present
	if len(s3PEM) > 0 {
		content += "\n" + string(s3PEM)
	}

	// Create directory if needed
	if idx := strings.LastIndex(path, "/"); idx > 0 {
		dir := path[:idx]
		if err := os.MkdirAll(dir, 0700); err != nil {
			return fmt.Errorf("failed to create directory: %w", err)
		}
	}

	// Write with restrictive permissions
	if err := os.WriteFile(path, []byte(content), 0600); err != nil {
		return fmt.Errorf("failed to write key file: %w", err)
	}

	return nil
}

// Load reads a backup key file from disk.
func Load(path string) (*BackupKeyFile, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("failed to read key file: %w", err)
	}

	return Parse(data)
}

// Parse parses backup key file content.
func Parse(data []byte) (*BackupKeyFile, error) {
	var privateKey *ecdsa.PrivateKey
	var certificate *x509.Certificate
	var caCert *x509.Certificate
	var s3Creds s3CredentialsJSON

	// Parse all PEM blocks
	rest := data
	for {
		var block *pem.Block
		block, rest = pem.Decode(rest)
		if block == nil {
			break
		}

		switch block.Type {
		case "EC PRIVATE KEY":
			key, err := x509.ParseECPrivateKey(block.Bytes)
			if err != nil {
				return nil, fmt.Errorf("failed to parse EC private key: %w", err)
			}
			privateKey = key

		case "CERTIFICATE":
			cert, err := x509.ParseCertificate(block.Bytes)
			if err != nil {
				return nil, fmt.Errorf("failed to parse certificate: %w", err)
			}

			// Check if this is the CA cert
			if block.Headers["X-Purpose"] == "CA" || cert.IsCA {
				caCert = cert
			} else if certificate == nil {
				// First non-CA cert is the client cert
				certificate = cert
			}

		case "PRIVATE KEY":
			// Handle PKCS#8 format
			key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
			if err != nil {
				return nil, fmt.Errorf("failed to parse PKCS8 private key: %w", err)
			}
			ecKey, ok := key.(*ecdsa.PrivateKey)
			if !ok {
				return nil, errors.New("private key is not ECDSA")
			}
			privateKey = ecKey

		case "PRIVATEBACKUP S3 CREDENTIALS":
			if err := json.Unmarshal(block.Bytes, &s3Creds); err != nil {
				return nil, fmt.Errorf("failed to parse S3 credentials: %w", err)
			}
		}
	}

	// Validate required components
	if privateKey == nil {
		return nil, errors.New("backup key file missing private key")
	}
	if certificate == nil {
		return nil, errors.New("backup key file missing client certificate")
	}
	if caCert == nil {
		return nil, errors.New("backup key file missing CA certificate")
	}

	// Verify certificate matches private key
	if !privateKeyMatchesCert(privateKey, certificate) {
		return nil, errors.New("private key does not match certificate")
	}

	return &BackupKeyFile{
		PrivateKey:        privateKey,
		Certificate:       certificate,
		CACert:            caCert,
		CertHash:          ComputeCertHash(certificate),
		S3AccessKeyID:     s3Creds.AccessKeyID,
		S3SecretAccessKey: s3Creds.SecretAccessKey,
		S3BucketName:      s3Creds.BucketName,
		CreatedAt:         time.Now(), // We don't preserve original creation time
	}, nil
}

// privateKeyMatchesCert verifies that a private key corresponds to a certificate.
func privateKeyMatchesCert(key *ecdsa.PrivateKey, cert *x509.Certificate) bool {
	pubKey, ok := cert.PublicKey.(*ecdsa.PublicKey)
	if !ok {
		return false
	}
	return key.PublicKey.Equal(pubKey)
}

// VerifyCertificate verifies that the client certificate was signed by the CA.
func (bkf *BackupKeyFile) VerifyCertificate() error {
	roots := x509.NewCertPool()
	roots.AddCert(bkf.CACert)

	opts := x509.VerifyOptions{
		Roots: roots,
		// Skip time validation since certs don't expire
		CurrentTime: bkf.Certificate.NotBefore.Add(time.Hour),
		// Specify key usages to validate
		KeyUsages: []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
	}

	if _, err := bkf.Certificate.Verify(opts); err != nil {
		return fmt.Errorf("certificate verification failed: %w", err)
	}

	return nil
}

// TLSCertificate returns a crypto/tls compatible certificate for mTLS.
func (bkf *BackupKeyFile) TLSCertificate() (cert []byte, key crypto.PrivateKey, err error) {
	return bkf.Certificate.Raw, bkf.PrivateKey, nil
}

// CreateCombinedPEM creates a combined PEM file containing both the client
// certificate and private key, suitable for use with restic's --tls-client-cert
// or RESTIC_TLS_CLIENT_CERT environment variable.
func (bkf *BackupKeyFile) CreateCombinedPEM(path string) error {
	// Encode certificate first (restic expects cert then key)
	certPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE",
		Bytes: bkf.Certificate.Raw,
	})

	// Encode private key
	keyBytes, err := x509.MarshalECPrivateKey(bkf.PrivateKey)
	if err != nil {
		return fmt.Errorf("failed to marshal private key: %w", err)
	}
	keyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "EC PRIVATE KEY",
		Bytes: keyBytes,
	})

	// Combine cert and key
	combined := append(certPEM, keyPEM...)

	// Create directory if needed
	if idx := strings.LastIndex(path, "/"); idx > 0 {
		dir := path[:idx]
		if err := os.MkdirAll(dir, 0700); err != nil {
			return fmt.Errorf("failed to create directory: %w", err)
		}
	}

	// Write with restrictive permissions
	if err := os.WriteFile(path, combined, 0600); err != nil {
		return fmt.Errorf("failed to write combined PEM file: %w", err)
	}

	return nil
}

// CreateCACertPEM creates a PEM file containing the CA certificate,
// suitable for use with restic's --cacert or RESTIC_CACERT environment variable.
func (bkf *BackupKeyFile) CreateCACertPEM(path string) error {
	caPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE",
		Bytes: bkf.CACert.Raw,
	})

	// Create directory if needed
	if idx := strings.LastIndex(path, "/"); idx > 0 {
		dir := path[:idx]
		if err := os.MkdirAll(dir, 0700); err != nil {
			return fmt.Errorf("failed to create directory: %w", err)
		}
	}

	// Write with restrictive permissions
	if err := os.WriteFile(path, caPEM, 0600); err != nil {
		return fmt.Errorf("failed to write CA PEM file: %w", err)
	}

	return nil
}

// GetRepositoryPassword derives and returns the Restic repository password.
func (bkf *BackupKeyFile) GetRepositoryPassword() (string, error) {
	return DeriveRepositoryPassword(bkf.Certificate, bkf.PrivateKey)
}

// Load loads the backup key file using the manager's configured path.
func (m *Manager) Load() (*BackupKeyFile, error) {
	return Load(m.keyFilePath)
}

// Save saves a backup key file using the manager's configured path.
func (m *Manager) Save(bkf *BackupKeyFile) error {
	return bkf.Save(m.keyFilePath)
}
