package verify

import (
	"crypto/ecdsa"
	"crypto/ed25519"
	"crypto/elliptic"
	"crypto/sha256"
	"encoding/base64"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/Fuwn/plutia/internal/config"
	"github.com/Fuwn/plutia/internal/types"
	"github.com/decred/dcrd/dcrec/secp256k1/v4"
	secpECDSA "github.com/decred/dcrd/dcrec/secp256k1/v4/ecdsa"
	"github.com/fxamacker/cbor/v2"
	"github.com/mr-tron/base58"
	"math/big"
	"strings"
)

type Verifier struct {
	Policy string
}

func New(policy string) *Verifier {
	return &Verifier{Policy: policy}
}

func (v *Verifier) ShouldVerify(existing *types.StateV1, seq uint64) bool {
	switch v.Policy {
	case config.VerifyLazy:
		return false
	case config.VerifyStateOnly:
		if existing == nil {
			return true
		}

		return seq >= existing.LatestOpSeq
	default:
		return true
	}
}

func (v *Verifier) VerifyOperation(op types.ParsedOperation, existing *types.StateV1) error {
	if !v.ShouldVerify(existing, op.Sequence) {
		return nil
	}

	if existing != nil && existing.ChainTipHash != "" {
		if op.Prev == "" {
			return errors.New("missing prev on non-genesis operation")
		}

		if op.Prev != existing.ChainTipHash {
			return fmt.Errorf("prev linkage mismatch: got %s want %s", op.Prev, existing.ChainTipHash)
		}
	}

	sig, ok := findString(op.Payload, "sig", "signature")

	if !ok || strings.TrimSpace(sig) == "" {
		return errors.New("missing signature")
	}

	pubKeys := make([]string, 0, 8)

	if existing != nil && len(existing.RotationKeys) > 0 {
		pubKeys = append(pubKeys, existing.RotationKeys...)
	}

	pubKeys = append(pubKeys, extractPublicKeys(op.Payload)...)

	if len(pubKeys) == 0 {
		return errors.New("missing verification public key")
	}

	sigBytes, err := decodeFlexible(sig)

	if err != nil {
		return fmt.Errorf("decode signature: %w", err)
	}

	payload, err := signaturePayload(op.Payload)

	if err != nil {
		return err
	}

	for _, key := range pubKeys {
		vk, err := decodePublicKey(key)

		if err != nil {
			continue
		}

		switch vk.Algo {
		case "ed25519":
			if ed25519.Verify(vk.Ed25519, payload, sigBytes) {
				return nil
			}
		case "secp256k1":
			if verifySecp256k1(vk.Secp256k1, payload, sigBytes) {
				return nil
			}
		case "p256":
			if verifyP256(vk.P256, payload, sigBytes) {
				return nil
			}
		}
	}

	return errors.New("signature verification failed")
}

func signaturePayload(m map[string]any) ([]byte, error) {
	if raw, ok := findString(m, "sigPayload", "signaturePayload"); ok && raw != "" {
		decoded, err := decodeFlexible(raw)

		if err == nil && json.Valid(decoded) {
			return types.CanonicalizeJSON(decoded)
		}
	}

	clone := make(map[string]any, len(m))

	for k, v := range m {
		switch k {
		case "sig", "signature", "sigPayload", "signaturePayload":
			continue
		default:
			clone[k] = v
		}
	}

	encMode, err := cbor.CanonicalEncOptions().EncMode()

	if err != nil {
		return nil, fmt.Errorf("init canonical cbor encoder: %w", err)
	}

	b, err := encMode.Marshal(clone)

	if err != nil {
		return nil, fmt.Errorf("marshal signature payload cbor: %w", err)
	}

	return b, nil
}

func findString(m map[string]any, keys ...string) (string, bool) {
	for _, k := range keys {
		if v, ok := m[k].(string); ok {
			return v, true
		}
	}

	return "", false
}

func extractPublicKeys(payload map[string]any) []string {
	out := make([]string, 0, 6)
	seen := map[string]struct{}{}
	add := func(v string) {
		v = strings.TrimSpace(v)

		if v == "" {
			return
		}

		if _, ok := seen[v]; ok {
			return
		}

		seen[v] = struct{}{}
		out = append(out, v)
	}

	if arr, ok := payload["rotationKeys"].([]any); ok {
		for _, v := range arr {
			if s, ok := v.(string); ok {
				add(s)
			}
		}
	}

	if v, ok := findString(payload, "publicKey", "verificationMethod", "signingKey", "recoveryKey"); ok {
		add(v)
	}

	if vm, ok := payload["verificationMethods"].(map[string]any); ok {
		if v, ok := vm["atproto"].(string); ok {
			add(v)
		}

		for _, anyV := range vm {
			if s, ok := anyV.(string); ok {
				add(s)
			}
		}
	}

	return out
}

type verificationKey struct {
	Algo      string
	Ed25519   ed25519.PublicKey
	Secp256k1 *secp256k1.PublicKey
	P256      *ecdsa.PublicKey
}

func decodePublicKey(value string) (verificationKey, error) {
	if strings.HasPrefix(value, "did:key:") {
		mb := strings.TrimPrefix(value, "did:key:")

		if mb == "" || mb[0] != 'z' {
			return verificationKey{}, errors.New("did:key must be multibase base58btc")
		}

		decoded, err := base58.Decode(mb[1:])

		if err != nil {
			return verificationKey{}, fmt.Errorf("decode did:key base58: %w", err)
		}

		if len(decoded) < 3 {
			return verificationKey{}, errors.New("invalid did:key length")
		}

		code, n := binary.Uvarint(decoded)

		if n <= 0 || n >= len(decoded) {
			return verificationKey{}, errors.New("invalid did:key multicodec prefix")
		}

		keyBytes := decoded[n:]

		switch code {
		case 0xED:
			if len(keyBytes) != ed25519.PublicKeySize {
				return verificationKey{}, errors.New("invalid did:key ed25519 length")
			}

			return verificationKey{
				Algo:    "ed25519",
				Ed25519: ed25519.PublicKey(keyBytes),
			}, nil
		case 0xE7:
			pub, err := secp256k1.ParsePubKey(keyBytes)

			if err != nil {
				return verificationKey{}, fmt.Errorf("parse secp256k1 did:key: %w", err)
			}

			return verificationKey{Algo: "secp256k1", Secp256k1: pub}, nil
		case 0x1200:
			x, y := elliptic.UnmarshalCompressed(elliptic.P256(), keyBytes)

			if x == nil || y == nil {
				return verificationKey{}, errors.New("parse p256 did:key: invalid compressed key")
			}

			return verificationKey{
				Algo: "p256",
				P256: &ecdsa.PublicKey{Curve: elliptic.P256(), X: x, Y: y},
			}, nil
		default:
			return verificationKey{}, errors.New("unsupported did:key multicodec")
		}
	}

	b, err := decodeFlexible(value)

	if err != nil {
		return verificationKey{}, fmt.Errorf("decode public key: %w", err)
	}

	if len(b) == ed25519.PublicKeySize {
		return verificationKey{Algo: "ed25519", Ed25519: ed25519.PublicKey(b)}, nil
	}

	pub, err := secp256k1.ParsePubKey(b)

	if err == nil {
		return verificationKey{Algo: "secp256k1", Secp256k1: pub}, nil
	}

	if x, y := elliptic.UnmarshalCompressed(elliptic.P256(), b); x != nil && y != nil {
		return verificationKey{
			Algo: "p256",
			P256: &ecdsa.PublicKey{Curve: elliptic.P256(), X: x, Y: y},
		}, nil
	}

	return verificationKey{}, fmt.Errorf("invalid public key length/type: %d", len(b))
}

func verifySecp256k1(pub *secp256k1.PublicKey, payload, sig []byte) bool {
	if pub == nil {
		return false
	}

	var parsed *secpECDSA.Signature

	if len(sig) == 64 {
		var r, s secp256k1.ModNScalar

		r.SetByteSlice(sig[:32])
		s.SetByteSlice(sig[32:])

		parsed = secpECDSA.NewSignature(&r, &s)
	} else {
		der, err := secpECDSA.ParseDERSignature(sig)

		if err != nil {
			return false
		}

		parsed = der
	}

	sum := sha256.Sum256(payload)

	return parsed.Verify(sum[:], pub)
}

func verifyP256(pub *ecdsa.PublicKey, payload, sig []byte) bool {
	if pub == nil {
		return false
	}

	sum := sha256.Sum256(payload)

	if len(sig) == 64 {
		r := new(big.Int).SetBytes(sig[:32])
		s := new(big.Int).SetBytes(sig[32:])

		return ecdsa.Verify(pub, sum[:], r, s)
	}

	return ecdsa.VerifyASN1(pub, sum[:], sig)
}

func decodeFlexible(v string) ([]byte, error) {
	if b, err := base64.RawURLEncoding.DecodeString(v); err == nil {
		return b, nil
	}

	if b, err := base64.StdEncoding.DecodeString(v); err == nil {
		return b, nil
	}

	if b, err := hex.DecodeString(v); err == nil {
		return b, nil
	}

	return nil, errors.New("unsupported encoding")
}