// Package stm implements Mithril Stake-based Threshold Multi-signature
// decoding and verification.
//
// The wire format of a Mithril multi_signature field is:
//
//	hex( ASCII( JSON( ... ) ) )
//
// i.e. hex-encoded bytes that are the UTF-8 of a JSON object. The JSON
// contents are documented in DecodeMultiSig below.
//
// Verification phases:
//
//  1. DecodeMultiSig — parse the wrapped JSON
//  2. BLS single-sig verification of each (signer, sigma) over the msg
//  3. Merkle proof verification: each signer index is a registered party
//  4. Lottery check: for each (index, sigma), evaluate_dense_mapping < threshold(stake)
//  5. Threshold: total distinct lottery wins >= k
//
// Phases 2-5 are implemented in verify.go (BLS, lottery, Merkle, threshold).
package stm

import (
	"encoding/hex"
	"encoding/json"
	"fmt"
)

// ByteArray decodes from either a JSON array of ints [1,2,3] (Mithril's
// on-wire shape) or a base64 string (Go's default []byte handling).
// Always emits an array of ints for forward compatibility.
type ByteArray []byte

func (b *ByteArray) UnmarshalJSON(data []byte) error {
	// Try array-of-ints first — this is what Mithril ships.
	var ints []int
	if err := json.Unmarshal(data, &ints); err == nil {
		out := make([]byte, len(ints))
		for i, v := range ints {
			if v < 0 || v > 255 {
				return fmt.Errorf("byte out of range at %d: %d", i, v)
			}
			out[i] = byte(v)
		}
		*b = out
		return nil
	}
	// Fallback: base64 string.
	var s string
	if err := json.Unmarshal(data, &s); err == nil {
		*b = []byte(s)
		return nil
	}
	return fmt.Errorf("ByteArray: neither int-array nor string")
}

// MultiSig is the decoded top-level shape.
type MultiSig struct {
	Signatures []SignerEntry `json:"signatures"`
	BatchProof BatchProof    `json:"batch_proof"`
}

// SignerEntry is a 2-tuple serialized as a JSON array: (StmSig, RegParty).
// We decode it via a custom UnmarshalJSON because JSON heterogeneous
// arrays don't map to Go structs directly.
type SignerEntry struct {
	Sig      StmSig
	RegParty RegParty
}

// StmSig is one signer's contribution: their BLS sig, the lottery
// indices they won, and their index in the registered party list.
type StmSig struct {
	Sigma       ByteArray `json:"sigma"`        // 48-byte BLS G1 compressed sig
	Indexes     []uint64  `json:"indexes"`      // winning lottery indices
	SignerIndex uint64    `json:"signer_index"` // party position in registered list
}

// RegParty is also a 2-tuple in JSON: (vk_bytes, stake).
type RegParty struct {
	VK    ByteArray // 96-byte BLS G2 compressed verification key
	Stake uint64
}

// AVK is the cert's aggregate verification key — a Merkle commitment over
// the registered (vk, stake) parties plus the total stake. Shipped on the
// wire as hex-of-ASCII-of-JSON, same wrapping as MultiSig.
type AVK struct {
	MerkleRoot ByteArray
	NumLeaves  uint64
	TotalStake uint64
}

// DecodeAVK decodes the wrapped JSON of a Mithril aggregate_verification_key
// field.
func DecodeAVK(rawJSON []byte) (*AVK, error) {
	hexStr := string(rawJSON)
	if len(hexStr) >= 2 && hexStr[0] == '"' && hexStr[len(hexStr)-1] == '"' {
		hexStr = hexStr[1 : len(hexStr)-1]
	}
	data, err := hex.DecodeString(hexStr)
	if err != nil {
		return nil, fmt.Errorf("AVK hex: %w", err)
	}
	var wire struct {
		MTCommitment struct {
			Root      ByteArray `json:"root"`
			NrLeaves  uint64    `json:"nr_leaves"`
			Hasher    any       `json:"hasher"`
		} `json:"mt_commitment"`
		TotalStake uint64 `json:"total_stake"`
	}
	if err := json.Unmarshal(data, &wire); err != nil {
		return nil, fmt.Errorf("AVK json: %w", err)
	}
	if len(wire.MTCommitment.Root) != 32 {
		return nil, fmt.Errorf("AVK root: got %d bytes, want 32", len(wire.MTCommitment.Root))
	}
	if wire.TotalStake == 0 {
		return nil, fmt.Errorf("AVK total_stake is zero")
	}
	if wire.MTCommitment.NrLeaves == 0 {
		return nil, fmt.Errorf("AVK nr_leaves is zero")
	}
	return &AVK{
		MerkleRoot: wire.MTCommitment.Root,
		NumLeaves:  wire.MTCommitment.NrLeaves,
		TotalStake: wire.TotalStake,
	}, nil
}

// BatchProof is a Merkle multi-proof over the registered parties.
type BatchProof struct {
	Values  []ByteArray `json:"values"`  // proof nodes, each 32 bytes (BLAKE2b-256)
	Indices []uint64    `json:"indices"` // signer indices being proven
	Hasher  any         `json:"hasher"`  // null => BLAKE2b-256 default
}

// UnmarshalJSON for SignerEntry — decode the [StmSig, RegParty] tuple.
func (s *SignerEntry) UnmarshalJSON(b []byte) error {
	var raw [2]json.RawMessage
	if err := json.Unmarshal(b, &raw); err != nil {
		return fmt.Errorf("SignerEntry tuple: %w", err)
	}
	if err := json.Unmarshal(raw[0], &s.Sig); err != nil {
		return fmt.Errorf("SignerEntry.Sig: %w", err)
	}
	if err := s.RegParty.UnmarshalJSON(raw[1]); err != nil {
		return fmt.Errorf("SignerEntry.RegParty: %w", err)
	}
	return nil
}

// UnmarshalJSON for RegParty — decode the [vk_bytes, stake] tuple.
func (r *RegParty) UnmarshalJSON(b []byte) error {
	var raw [2]json.RawMessage
	if err := json.Unmarshal(b, &raw); err != nil {
		return fmt.Errorf("RegParty tuple: %w", err)
	}
	if err := r.VK.UnmarshalJSON(raw[0]); err != nil {
		return fmt.Errorf("RegParty.VK: %w", err)
	}
	if err := json.Unmarshal(raw[1], &r.Stake); err != nil {
		return fmt.Errorf("RegParty.Stake: %w", err)
	}
	return nil
}

// DecodeMultiSig takes the raw `multi_signature` field value from a
// Mithril certificate (a JSON string whose contents are hex-encoded
// UTF-8 JSON) and returns the decoded struct.
//
// If rawJSON begins with a JSON string quote, the quotes are stripped
// first; this lets callers pass either the json.RawMessage form or an
// already-unquoted hex string.
func DecodeMultiSig(rawJSON []byte) (*MultiSig, error) {
	hexStr := string(rawJSON)
	if len(hexStr) >= 2 && hexStr[0] == '"' && hexStr[len(hexStr)-1] == '"' {
		hexStr = hexStr[1 : len(hexStr)-1]
	}
	data, err := hex.DecodeString(hexStr)
	if err != nil {
		return nil, fmt.Errorf("hex decode: %w", err)
	}
	var ms MultiSig
	if err := json.Unmarshal(data, &ms); err != nil {
		return nil, fmt.Errorf("json decode: %w", err)
	}
	return &ms, nil
}

// TotalWins counts the total number of lottery wins across all signers.
func (m *MultiSig) TotalWins() int {
	n := 0
	for _, s := range m.Signatures {
		n += len(s.Sig.Indexes)
	}
	return n
}

// DistinctWins returns the set of distinct lottery indices claimed
// across all signers. The Mithril STM spec requires total DISTINCT
// indices >= k for a valid aggregate.
func (m *MultiSig) DistinctWins() []uint64 {
	seen := make(map[uint64]struct{})
	for _, s := range m.Signatures {
		for _, ix := range s.Sig.Indexes {
			seen[ix] = struct{}{}
		}
	}
	out := make([]uint64, 0, len(seen))
	for ix := range seen {
		out = append(out, ix)
	}
	return out
}