package authsnapshot

import (
	"bytes"
	"crypto/ed25519"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"sort"
	"strconv"
	"strings"
	"time"
)

const (
	SnapshotVersion = 1
	MaxTTL          = 5 * time.Minute
	clockSkew       = 30 * time.Second
	signatureDomain = "nekonest-cloud/relay-authorization-snapshot/v1\n"
)

type Device struct {
	DeviceID            string `json:"device_id"`
	CredentialHash      string `json:"credential_hash"`
	IdentityFingerprint string `json:"identity_fingerprint"`
	Ed25519Public       string `json:"ed25519_public"`
	X25519Public        string `json:"x25519_public"`
	Name                string `json:"name"`
	OS                  string `json:"os"`
}

type Phone struct {
	PhoneID             string `json:"phone_id"`
	CredentialHash      string `json:"credential_hash"`
	IdentityFingerprint string `json:"identity_fingerprint"`
	Ed25519Public       string `json:"ed25519_public"`
	X25519Public        string `json:"x25519_public"`
	Name                string `json:"name"`
}

type Payload struct {
	SnapshotVersion       int      `json:"snapshot_version"`
	TenantID              string   `json:"tenant_id"`
	TenantStatus          string   `json:"tenant_status"`
	HomeRegion            string   `json:"home_region"`
	RelayNodeID           string   `json:"relay_node_id"`
	PlacementGeneration   int64    `json:"placement_generation"`
	AuthorizationRevision int64    `json:"authorization_revision"`
	Devices               []Device `json:"devices"`
	Phones                []Phone  `json:"phones,omitempty"`
	IssuedAt              string   `json:"issued_at"`
	ExpiresAt             string   `json:"expires_at"`
}

type Signed struct {
	Algorithm string  `json:"algorithm"`
	KID       string  `json:"kid"`
	Payload   Payload `json:"payload"`
	Signature string  `json:"signature"`
}

type Key struct {
	PublicKey   ed25519.PublicKey
	NotBefore   time.Time
	RetainUntil time.Time
}

type Keyring map[string]Key

func parseTime(value string) (time.Time, error) {
	parsed, err := time.Parse(time.RFC3339Nano, value)
	if err != nil {
		return time.Time{}, fmt.Errorf("invalid snapshot time: %w", err)
	}
	return parsed, nil
}

func validPrefixedHex(value, prefix string, hexLength int) bool {
	if !strings.HasPrefix(value, prefix) || len(value) != len(prefix)+hexLength {
		return false
	}
	for _, char := range strings.TrimPrefix(value, prefix) {
		if !strings.ContainsRune("0123456789abcdef", char) {
			return false
		}
	}
	return true
}

func validDigest(value string) bool {
	if len(value) != 64 {
		return false
	}
	for _, char := range value {
		if !strings.ContainsRune("0123456789abcdef", char) {
			return false
		}
	}
	return true
}

func validPrefixedID(value, prefix string, maxSuffix int) bool {
	if !strings.HasPrefix(value, prefix) || len(value) <= len(prefix) || len(value) > len(prefix)+maxSuffix {
		return false
	}
	for _, char := range strings.TrimPrefix(value, prefix) {
		if !(char >= 'A' && char <= 'Z') && !(char >= 'a' && char <= 'z') &&
			!(char >= '0' && char <= '9') && !strings.ContainsRune("._:-", char) {
			return false
		}
	}
	return true
}

func validBase64URLKey(value string) bool {
	if len(value) != 43 {
		return false
	}
	for _, char := range value {
		if !(char >= 'A' && char <= 'Z') && !(char >= 'a' && char <= 'z') &&
			!(char >= '0' && char <= '9') && char != '_' && char != '-' {
			return false
		}
	}
	return true
}

func validatePayload(payload Payload, now time.Time) (time.Time, time.Time, error) {
	if payload.SnapshotVersion != SnapshotVersion || !validPrefixedHex(payload.TenantID, "tenant_", 32) {
		return time.Time{}, time.Time{}, errors.New("invalid snapshot identity")
	}
	if !strings.HasPrefix(payload.RelayNodeID, "node_") || len(payload.RelayNodeID) > 101 {
		return time.Time{}, time.Time{}, errors.New("invalid relay node")
	}
	if payload.TenantStatus != "active" || payload.PlacementGeneration < 1 || payload.AuthorizationRevision < 0 {
		return time.Time{}, time.Time{}, errors.New("snapshot does not authorize an active placement")
	}
	issuedAt, err := parseTime(payload.IssuedAt)
	if err != nil {
		return time.Time{}, time.Time{}, err
	}
	expiresAt, err := parseTime(payload.ExpiresAt)
	if err != nil {
		return time.Time{}, time.Time{}, err
	}
	if !expiresAt.After(issuedAt) || expiresAt.Sub(issuedAt) > MaxTTL {
		return time.Time{}, time.Time{}, errors.New("invalid snapshot lifetime")
	}
	if now.Before(issuedAt.Add(-clockSkew)) || !now.Before(expiresAt) {
		return time.Time{}, time.Time{}, errors.New("snapshot is not current")
	}
	previous := ""
	for _, device := range payload.Devices {
		if !validPrefixedHex(device.DeviceID, "host_", 32) || !validDigest(device.CredentialHash) {
			return time.Time{}, time.Time{}, errors.New("invalid snapshot device")
		}
		if strings.TrimSpace(device.Name) == "" || len(device.Name) > 48 ||
			(device.OS != "windows" && device.OS != "linux") ||
			!validBase64URLKey(device.Ed25519Public) || !validBase64URLKey(device.X25519Public) ||
			!validDigest(device.IdentityFingerprint) {
			return time.Time{}, time.Time{}, errors.New("incomplete snapshot device identity")
		}
		if previous != "" && device.DeviceID <= previous {
			return time.Time{}, time.Time{}, errors.New("snapshot devices are not strictly ordered")
		}
		previous = device.DeviceID
	}
	previous = ""
	for _, phone := range payload.Phones {
		if !validPrefixedID(phone.PhoneID, "phone_", 120) || !validDigest(phone.CredentialHash) ||
			strings.TrimSpace(phone.Name) == "" || len(phone.Name) > 48 ||
			!validBase64URLKey(phone.Ed25519Public) || !validBase64URLKey(phone.X25519Public) ||
			!validDigest(phone.IdentityFingerprint) {
			return time.Time{}, time.Time{}, errors.New("invalid snapshot phone")
		}
		if previous != "" && phone.PhoneID <= previous {
			return time.Time{}, time.Time{}, errors.New("snapshot phones are not strictly ordered")
		}
		previous = phone.PhoneID
	}
	return issuedAt, expiresAt, nil
}

func canonicalJSON(value any) ([]byte, error) {
	encoded, err := json.Marshal(value)
	if err != nil {
		return nil, err
	}
	decoder := json.NewDecoder(bytes.NewReader(encoded))
	decoder.UseNumber()
	var decoded any
	if err := decoder.Decode(&decoded); err != nil {
		return nil, err
	}
	var output bytes.Buffer
	if err := writeCanonical(&output, decoded); err != nil {
		return nil, err
	}
	return output.Bytes(), nil
}

func writeCanonical(output *bytes.Buffer, value any) error {
	switch typed := value.(type) {
	case nil:
		output.WriteString("null")
	case bool:
		output.WriteString(strconv.FormatBool(typed))
	case string:
		encoded, _ := json.Marshal(typed)
		output.Write(encoded)
	case json.Number:
		if _, err := strconv.ParseInt(string(typed), 10, 64); err != nil {
			return fmt.Errorf("snapshot contains unsupported number %q", typed)
		}
		output.WriteString(string(typed))
	case []any:
		output.WriteByte('[')
		for index, item := range typed {
			if index > 0 {
				output.WriteByte(',')
			}
			if err := writeCanonical(output, item); err != nil {
				return err
			}
		}
		output.WriteByte(']')
	case map[string]any:
		keys := make([]string, 0, len(typed))
		for key := range typed {
			keys = append(keys, key)
		}
		sort.Strings(keys)
		output.WriteByte('{')
		for index, key := range keys {
			if index > 0 {
				output.WriteByte(',')
			}
			encoded, _ := json.Marshal(key)
			output.Write(encoded)
			output.WriteByte(':')
			if err := writeCanonical(output, typed[key]); err != nil {
				return err
			}
		}
		output.WriteByte('}')
	default:
		return fmt.Errorf("unsupported snapshot JSON value %T", value)
	}
	return nil
}

func SigningBytes(payload Payload) ([]byte, error) {
	canonical, err := canonicalJSON(payload)
	if err != nil {
		return nil, err
	}
	return append([]byte(signatureDomain), canonical...), nil
}

func Verify(snapshot Signed, keyring Keyring, now time.Time) (Payload, error) {
	if snapshot.Algorithm != "Ed25519" || snapshot.KID == "" {
		return Payload{}, errors.New("unsupported snapshot signature")
	}
	issuedAt, expiresAt, err := validatePayload(snapshot.Payload, now)
	if err != nil {
		return Payload{}, err
	}
	key, ok := keyring[snapshot.KID]
	if !ok || len(key.PublicKey) != ed25519.PublicKeySize {
		return Payload{}, errors.New("snapshot signing key is not pinned")
	}
	if issuedAt.Before(key.NotBefore) || key.RetainUntil.Before(expiresAt) {
		return Payload{}, errors.New("snapshot signing key lifetime is insufficient")
	}
	signature, err := base64.RawURLEncoding.DecodeString(snapshot.Signature)
	if err != nil || len(signature) != ed25519.SignatureSize {
		return Payload{}, errors.New("invalid snapshot signature encoding")
	}
	message, err := SigningBytes(snapshot.Payload)
	if err != nil {
		return Payload{}, err
	}
	if !ed25519.Verify(key.PublicKey, message, signature) {
		return Payload{}, errors.New("snapshot signature mismatch")
	}
	return snapshot.Payload, nil
}

func PublicKeyFromJWK(raw json.RawMessage) (ed25519.PublicKey, string, error) {
	var jwk struct {
		KTY string `json:"kty"`
		CRV string `json:"crv"`
		X   string `json:"x"`
		KID string `json:"kid"`
	}
	if err := json.Unmarshal(raw, &jwk); err != nil || jwk.KTY != "OKP" || jwk.CRV != "Ed25519" {
		return nil, "", errors.New("invalid Ed25519 JWK")
	}
	key, err := base64.RawURLEncoding.DecodeString(jwk.X)
	if err != nil || len(key) != ed25519.PublicKeySize {
		return nil, "", errors.New("invalid Ed25519 JWK key")
	}
	return ed25519.PublicKey(key), jwk.KID, nil
}