// Package verify implements signature verification for Mithril certificates.
//
// Two layers of verification exist in Mithril:
//
//  1. The genesis certificate is signed by a static Ed25519 key baked into
//     the client (per-network). This bootstraps trust into the STM protocol.
//  2. Subsequent certificates carry an STM (Stake-based Threshold Multi-
//     signature) aggregate signature over BLS12-381.
//
// Ed25519 (genesis) verification is fully implemented here. STM verification
// is stubbed pending the BLS crypto sprint.
package verify

import (
	"crypto/ed25519"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"sort"
	"strconv"
	"strings"
)

var (
	ErrNotGenesis        = errors.New("certificate is not a genesis certificate")
	ErrBadSignature      = errors.New("genesis signature verification failed")
	ErrSignedMessageHash = errors.New("signed_message does not match SHA256(protocol_message)")
)

// The Mithril enum order on ProtocolMessagePartKey — BTreeMap iteration
// follows variant declaration order, which is what compute_hash relies on.
// Keep this aligned with mithril-common/src/entities/protocol_message.rs.
var messagePartOrder = []string{
	"snapshot_digest",
	"cardano_transactions_merkle_root",
	"cardano_blocks_transactions_merkle_root",
	"next_aggregate_verification_key",
	"next_protocol_parameters",
	"current_epoch",
	"latest_block_number",
	"cardano_blocks_transactions_block_number_offset",
	"cardano_stake_distribution_epoch",
	"cardano_stake_distribution_merkle_root",
	"cardano_database_merkle_root",
	"next_aggregate_verification_key_snark",
}

// ComputeProtocolMessageHash mirrors Rust's ProtocolMessage::compute_hash:
// SHA256 over each (key_string || value_string) concatenated in
// BTreeMap order (Rust enum declaration order, not alphabetical).
// Returns the lowercase-hex form so it can be compared to the
// signed_message field directly.
func ComputeProtocolMessageHash(messageParts map[string]string) string {
	h := sha256.New()
	// iterate in canonical order
	ordinal := make(map[string]int, len(messagePartOrder))
	for i, k := range messagePartOrder {
		ordinal[k] = i
	}
	keys := make([]string, 0, len(messageParts))
	for k := range messageParts {
		if _, known := ordinal[k]; known {
			keys = append(keys, k)
		} else {
			// Unknown key — placed at end, alphabetical. Safer than
			// silently dropping; any future key we haven't mirrored yet
			// still feeds into the digest, matching Rust's forward-
			// compatibility.
			keys = append(keys, k)
		}
	}
	sort.Slice(keys, func(i, j int) bool {
		oi, iOK := ordinal[keys[i]]
		oj, jOK := ordinal[keys[j]]
		switch {
		case iOK && jOK:
			return oi < oj
		case iOK:
			return true
		case jOK:
			return false
		default:
			return keys[i] < keys[j]
		}
	})
	for _, k := range keys {
		h.Write([]byte(k))
		h.Write([]byte(messageParts[k]))
	}
	return hex.EncodeToString(h.Sum(nil))
}

// ProtocolMessage is the subset of the Mithril protocol_message JSON shape
// that verification consumes.
type ProtocolMessage struct {
	MessageParts map[string]string `json:"message_parts"`
}

// DecodeGenesisVerifyKey unpacks the Mithril genesis verification key from
// its serialized form. The key is stored as:
//
//	hex( ascii( "[b0,b1,...,b31]" ) )
//
// i.e. the Ed25519 public key's 32 raw bytes are written as a Rust-debug
// byte-slice literal, converted to ASCII bytes, then hex-encoded. We
// strip both layers and return the 32-byte public key.
//
// Plain 64-char hex of the raw key is also accepted (some distributions
// ship the key that way).
func DecodeGenesisVerifyKey(encoded string) (ed25519.PublicKey, error) {
	// Fast-path: plain 64-char hex of 32 raw bytes.
	if len(encoded) == 2*ed25519.PublicKeySize {
		if b, err := hex.DecodeString(encoded); err == nil && len(b) == ed25519.PublicKeySize {
			return ed25519.PublicKey(b), nil
		}
	}
	// Standard path: hex-of-ASCII-of-"[a,b,c,...]".
	inner, err := hex.DecodeString(encoded)
	if err != nil {
		return nil, fmt.Errorf("decode outer hex: %w", err)
	}
	s := strings.TrimSpace(string(inner))
	s = strings.TrimPrefix(s, "[")
	s = strings.TrimSuffix(s, "]")
	parts := strings.Split(s, ",")
	if len(parts) != ed25519.PublicKeySize {
		return nil, fmt.Errorf("genesis key literal has %d parts, want %d", len(parts), ed25519.PublicKeySize)
	}
	out := make([]byte, ed25519.PublicKeySize)
	for i, p := range parts {
		n, err := strconv.Atoi(strings.TrimSpace(p))
		if err != nil {
			return nil, fmt.Errorf("parse byte %d: %w", i, err)
		}
		if n < 0 || n > 255 {
			return nil, fmt.Errorf("byte %d out of range: %d", i, n)
		}
		out[i] = byte(n)
	}
	return ed25519.PublicKey(out), nil
}

// Genesis verifies a genesis certificate:
//   - signedMessageHex must equal SHA256(protocolMessage) in hex
//   - the Ed25519 signature (64 bytes hex) over the ASCII bytes of
//     signedMessageHex must verify against verifyKey
//
// Returns nil on success, a typed error on any failure. Matches the
// upstream `verify_genesis_certificate` semantics (integrity →
// signed_message → signature → epoch-match; the epoch check is the
// caller's job since it needs the full cert context).
func Genesis(verifyKey ed25519.PublicKey, signedMessageHex, genesisSignatureHex string, pm ProtocolMessage) error {
	computed := ComputeProtocolMessageHash(pm.MessageParts)
	if computed != signedMessageHex {
		return fmt.Errorf("%w: got %s, want %s", ErrSignedMessageHash, computed, signedMessageHex)
	}
	sig, err := hex.DecodeString(genesisSignatureHex)
	if err != nil {
		return fmt.Errorf("decode signature: %w", err)
	}
	if len(sig) != ed25519.SignatureSize {
		return fmt.Errorf("signature wrong size: got %d, want %d", len(sig), ed25519.SignatureSize)
	}
	// Per upstream: verify over signed_message.as_bytes() — the ASCII
	// bytes of the hex string, not the 32 raw digest bytes.
	if !ed25519.Verify(verifyKey, []byte(signedMessageHex), sig) {
		return ErrBadSignature
	}
	return nil
}

// GenesisFromJSON is a convenience wrapper when the caller has the raw
// JSON protocol_message (as the aggregator returns).
func GenesisFromJSON(verifyKey ed25519.PublicKey, signedMessageHex, genesisSignatureHex string, protocolMessageJSON []byte) error {
	var pm ProtocolMessage
	if err := json.Unmarshal(protocolMessageJSON, &pm); err != nil {
		return fmt.Errorf("parse protocol_message: %w", err)
	}
	return Genesis(verifyKey, signedMessageHex, genesisSignatureHex, pm)
}

// STM verification lives in the sibling internal/stm package — see
// stm.Verify(). This file is genesis-Ed25519-only.