package main

import (
	"context"
	"crypto/rand"
	"fmt"
	"time"

	"golang.org/x/crypto/nacl/box"
	"golang.org/x/crypto/nacl/secretbox"
)

const (
	fileKeyBytes        = 32
	secretboxNonceBytes = 24
	secretboxMACBytes   = 16
)

type FileKeyData struct {
	FileKey      [fileKeyBytes]byte
	NonceForData [secretboxNonceBytes]byte
}

type AppDataRecord struct {
	FileID        string
	Version       string
	DataPlain     []byte
	DataEncrypted []byte
}

func (c *Client) RequestFileKey(
	ctx context.Context,
	appID string,
	pollItv time.Duration,
) (*FileKeyData, string, error) {
	if pollItv <= 0 {
		return nil, "", fmt.Errorf("poll interval must be positive")
	}
	identity := "app." + appID
	for {
		body, err := c.PostKPS(identity, "app_req_filekey", map[string]any{})
		if err != nil {
			return nil, "", err
		}
		switch asString(body["status"]) {
		case "done":
			fileKey, err := c.OpenDataFileKey(toBytes(body["data_filekey"]))
			return fileKey, asString(body["ver"]), err
		case "denied":
			return nil, "", fmt.Errorf("Safe denied the filekey request")
		case "safe_auth_required":
			timer := time.NewTimer(pollItv)
			select {
			case <-ctx.Done():
				timer.Stop()
				return nil, "", ctx.Err()
			case <-timer.C:
			}
		default:
			return nil, "", fmt.Errorf("unexpected app_req_filekey status")
		}
	}
}

func (c *Client) OpenDataFileKey(blob []byte) (*FileKeyData, error) {
	envelope, err := unpackMap(blob)
	if err != nil {
		return nil, fmt.Errorf("decode data_filekey: %w", err)
	}
	encryptedKey := firstBytes(envelope, "encFileKey", "k")
	nonceBytes := firstBytes(envelope, "nonceForKey", "nk")
	legacyNonceBytes := firstBytes(envelope, "nonceForData", "nd")
	senderPkBytes := firstBytes(envelope, "apEncPk", "p")
	if len(encryptedKey) != fileKeyBytes+secretboxMACBytes ||
		len(nonceBytes) != secretboxNonceBytes ||
		len(legacyNonceBytes) != secretboxNonceBytes ||
		len(senderPkBytes) != fileKeyBytes {
		return nil, fmt.Errorf("data_filekey has invalid field lengths")
	}

	var nonce [secretboxNonceBytes]byte
	var senderPk [fileKeyBytes]byte
	copy(nonce[:], nonceBytes)
	copy(senderPk[:], senderPkBytes)
	opened, ok := box.Open(nil, encryptedKey, &nonce, &senderPk, c.encSk)
	if !ok || len(opened) != fileKeyBytes {
		return nil, fmt.Errorf("data_filekey box decryption failed")
	}

	fileKey := &FileKeyData{}
	copy(fileKey.FileKey[:], opened)
	copy(fileKey.NonceForData[:], legacyNonceBytes)
	return fileKey, nil
}

func (c *Client) OpenAccountKey(payload, safeEncPk []byte) ([]byte, error) {
	if len(payload) != secretboxNonceBytes+fileKeyBytes+box.Overhead ||
		len(safeEncPk) != fileKeyBytes {
		return nil, fmt.Errorf("Account Key envelope has invalid field lengths")
	}
	var nonce [secretboxNonceBytes]byte
	var senderPk [fileKeyBytes]byte
	copy(nonce[:], payload[:secretboxNonceBytes])
	copy(senderPk[:], safeEncPk)
	opened, ok := box.Open(
		nil,
		payload[secretboxNonceBytes:],
		&nonce,
		&senderPk,
		c.encSk,
	)
	if !ok || len(opened) != fileKeyBytes {
		return nil, fmt.Errorf("Account Key box decryption failed")
	}
	return opened, nil
}

func (c *Client) GetAppData(appID string) (*AppDataRecord, error) {
	body, err := c.PostKPS("app."+appID, "app_get_data", map[string]any{})
	if err != nil {
		return nil, err
	}
	record := &AppDataRecord{
		FileID:        asString(body["file_id"]),
		Version:       asString(body["ver"]),
		DataPlain:     toBytes(body["data_plain"]),
		DataEncrypted: toBytes(body["data_encrypted"]),
	}
	if record.FileID == "" || record.Version == "" {
		return nil, fmt.Errorf("app_get_data response is missing file_id or ver")
	}
	return record, nil
}

func (c *Client) SetAppData(
	appID string,
	version string,
	plaintext []byte,
	fileKey *FileKeyData,
) (string, error) {
	if version == "" {
		return "", fmt.Errorf("app_set_data requires a non-empty version")
	}
	ciphertext, err := EncryptAppData(plaintext, fileKey)
	if err != nil {
		return "", err
	}
	body, err := c.PostKPS("app."+appID, "app_set_data", map[string]any{
		"ver":      version,
		"data_enc": ciphertext,
	})
	if err != nil {
		return "", err
	}
	newVer := asString(body["ver"])
	if asString(body["file_id"]) == "" || newVer == "" {
		return "", fmt.Errorf("app_set_data response is missing file_id or ver")
	}
	return newVer, nil
}

func EncryptAppData(plaintext []byte, fileKey *FileKeyData) ([]byte, error) {
	if fileKey == nil {
		return nil, fmt.Errorf("filekey is required")
	}
	var nonce [secretboxNonceBytes]byte
	if _, err := rand.Read(nonce[:]); err != nil {
		return nil, err
	}
	out := make([]byte, 0, secretboxNonceBytes+secretboxMACBytes+len(plaintext))
	out = append(out, nonce[:]...)
	out = secretbox.Seal(out, plaintext, &nonce, &fileKey.FileKey)
	return out, nil
}

func DecryptAppData(ciphertext []byte, fileKey *FileKeyData) ([]byte, error) {
	if fileKey == nil {
		return nil, fmt.Errorf("filekey is required")
	}
	if len(ciphertext) >= secretboxNonceBytes+secretboxMACBytes {
		var nonce [secretboxNonceBytes]byte
		copy(nonce[:], ciphertext[:secretboxNonceBytes])
		if opened, ok := secretbox.Open(
			nil,
			ciphertext[secretboxNonceBytes:],
			&nonce,
			&fileKey.FileKey,
		); ok {
			return opened, nil
		}
	}
	if opened, ok := secretbox.Open(
		nil,
		ciphertext,
		&fileKey.NonceForData,
		&fileKey.FileKey,
	); ok {
		return opened, nil
	}
	return nil, fmt.Errorf("AppData decryption failed for v2 and legacy v1")
}

func firstBytes(values map[string]any, primary, legacy string) []byte {
	if value := toBytes(values[primary]); value != nil {
		return value
	}
	return toBytes(values[legacy])
}