package attesto

// Attesto 3 provenance verification (ATTESTO-PROVENANCE-001).
//
// Rust (edge/) is normative: it constructs commitments and capsule roots. This
// file is a verification client — it re-derives what the edge core produced and
// checks it. It cannot generate a randomizer, because outside the Local Vault
// there is nothing legitimate to commit.
//
// Canonicalization and domain hashing come from proofstream.go rather than a
// second implementation: the provenance lane hashes bytes under the same frozen
// ATTESTO-CANONICAL-JSON-001 rules as the rest of Attesto.
//
// Conformance is defined by golden-vectors/provenance-v0.1-dev/.

import (
	"crypto/subtle"
	"encoding/hex"
	"fmt"
	"sort"
)

const (
	ProvenanceProtocol        = "ATTESTO-PROVENANCE-001"
	ProvenanceProtocolVersion = "0.1"
	provenanceRandomizerBytes = 32
)

// ProvenanceDomains is the closed v1 registry. There is deliberately no generic
// attesto.provenance.v1.commitment fallback: every semantic object has its own
// domain, and an object without one is a protocol-registry change.
var ProvenanceDomains = map[string]struct{}{
	"attesto.provenance.v1.asset":           {},
	"attesto.provenance.v1.claim":           {},
	"attesto.provenance.v1.evidence":        {},
	"attesto.provenance.v1.edge":            {},
	"attesto.provenance.v1.capsule_leaf":    {},
	"attesto.provenance.v1.capsule_node":    {},
	"attesto.provenance.v1.capsule_root":    {},
	"attesto.provenance.v1.envelope":        {},
	"attesto.provenance.v1.provider_result": {},
	"attesto.provenance.v1.policy_result":   {},
	"attesto.provenance.v1.disclosure":      {},
	"attesto.provenance.v1.migration":       {},
	"attesto.provenance.v1.vault_identity":  {},
	"attesto.provenance.v1.attesto_mark":    {},
	"attesto.provenance.v1.claims_tree":     {},
	"attesto.provenance.v1.evidence_tree":   {},
	"attesto.provenance.v1.policy_tree":     {},
	"attesto.provenance.v1.attestation":     {},
}

// TopLeafRoles is the canonical order of the six typed top-tree leaves. No
// other leaf exists in v1.
var TopLeafRoles = [6]string{
	"subject_commitment",
	"claims_root",
	"evidence_root",
	"policy_results_root",
	"attestation_commitment",
	"vault_identity_commitment",
}

var subtreeTreeDomain = map[string]string{
	"claims":         "attesto.provenance.v1.claims_tree",
	"evidence":       "attesto.provenance.v1.evidence_tree",
	"policy_results": "attesto.provenance.v1.policy_tree",
}

var subtreeTopRole = map[string]string{
	"claims":         "claims_root",
	"evidence":       "evidence_root",
	"policy_results": "policy_results_root",
}

func assertProvenanceDomain(domain string) error {
	if _, ok := ProvenanceDomains[domain]; !ok {
		return fmt.Errorf("unknown provenance domain: %q; there is no fallback domain", domain)
	}
	return nil
}

func assertRandomizer(randomizer string) error {
	if len(randomizer) != provenanceRandomizerBytes*2 {
		return fmt.Errorf("randomizer must be exactly %d bytes", provenanceRandomizerBytes)
	}
	for _, char := range randomizer {
		if !((char >= '0' && char <= '9') || (char >= 'a' && char <= 'f')) {
			return fmt.Errorf("randomizer must be lowercase hex")
		}
	}
	if _, err := hex.DecodeString(randomizer); err != nil {
		return fmt.Errorf("randomizer must be lowercase hex")
	}
	return nil
}

func assertProvenanceDigest(field, digest string) error {
	if len(digest) != 64 {
		return fmt.Errorf("malformed digest in %s: expected 64 lowercase hex characters", field)
	}
	for _, char := range digest {
		if !((char >= '0' && char <= '9') || (char >= 'a' && char <= 'f')) {
			return fmt.Errorf("malformed digest in %s: expected 64 lowercase hex characters", field)
		}
	}
	return nil
}

// ProvenanceCommitmentDigest re-derives a randomized commitment. Verification
// only — the randomizer must already be known, which means the holder was given
// the opening.
func ProvenanceCommitmentDigest(domain string, value any, randomizer string) (string, error) {
	if err := assertProvenanceDomain(domain); err != nil {
		return "", err
	}
	if err := assertRandomizer(randomizer); err != nil {
		return "", err
	}
	if err := AssertCommitmentSafeNumbers(value, "$"); err != nil {
		return "", err
	}
	return DomainHashHex(domain, map[string]any{
		"protocol":         ProvenanceProtocol,
		"protocol_version": ProvenanceProtocolVersion,
		"randomizer":       randomizer,
		"value":            value,
	})
}

// VerifyProvenanceCommitment checks in constant time that (value, randomizer)
// opens expectedDigest.
func VerifyProvenanceCommitment(domain string, value any, randomizer, expectedDigest string) (bool, error) {
	digest, err := ProvenanceCommitmentDigest(domain, value, randomizer)
	if err != nil {
		return false, err
	}
	return subtle.ConstantTimeCompare([]byte(digest), []byte(expectedDigest)) == 1, nil
}

func provenanceNode(domain, left, right string) (string, error) {
	return DomainHashHex(domain, map[string]any{
		"kind":  "node",
		"left":  left,
		"right": right,
	})
}

func provenanceFold(domain string, level []string) (string, error) {
	current := append([]string(nil), level...)
	for len(current) > 1 {
		next := make([]string, 0, (len(current)+1)/2)
		for index := 0; index < len(current); index += 2 {
			if index+1 >= len(current) {
				next = append(next, current[index]) // promote odd node, never duplicate
				continue
			}
			node, err := provenanceNode(domain, current[index], current[index+1])
			if err != nil {
				return "", err
			}
			next = append(next, node)
		}
		current = next
	}
	return current[0], nil
}

// SubtreeMerkleRoot folds a subtree's ordered leaves into its bare Merkle root.
func SubtreeMerkleRoot(subtree string, orderedLeaves []string) (string, error) {
	domain, ok := subtreeTreeDomain[subtree]
	if !ok {
		return "", fmt.Errorf("unknown subtree: %q", subtree)
	}
	if len(orderedLeaves) == 0 {
		return "", fmt.Errorf("cannot build an empty %s tree", subtree)
	}
	for index, leaf := range orderedLeaves {
		if err := assertProvenanceDigest(fmt.Sprintf("orderedLeaves[%d]", index), leaf); err != nil {
			return "", err
		}
	}
	return provenanceFold(domain, orderedLeaves)
}

// SubtreeRoot wraps a bare Merkle root in its typed subtree root.
func SubtreeRoot(subtree, merkleRoot string, leafCount int) (string, error) {
	domain, ok := subtreeTreeDomain[subtree]
	if !ok {
		return "", fmt.Errorf("unknown subtree: %q", subtree)
	}
	if err := assertProvenanceDigest("merkleRoot", merkleRoot); err != nil {
		return "", err
	}
	return DomainHashHex(domain, map[string]any{
		"kind":        "root",
		"tree":        subtree,
		"leaf_count":  leafCount,
		"merkle_root": merkleRoot,
	})
}

// SubtreeLeafInput is one leaf awaiting canonical ordering.
type SubtreeLeafInput struct {
	LeafRole   string `json:"leaf_role"`
	LeafID     string `json:"leaf_id"`
	Commitment string `json:"commitment"`
}

// OrderSubtreeLeaves orders leaves by (leaf_role, leaf_id), the frozen rule, so
// two vaults that assembled the same facts in different orders agree.
func OrderSubtreeLeaves(leaves []SubtreeLeafInput) ([]string, error) {
	ordered := append([]SubtreeLeafInput(nil), leaves...)
	sort.SliceStable(ordered, func(left, right int) bool {
		if ordered[left].LeafRole != ordered[right].LeafRole {
			return ordered[left].LeafRole < ordered[right].LeafRole
		}
		return ordered[left].LeafID < ordered[right].LeafID
	})

	seen := make(map[string]struct{}, len(ordered))
	digests := make([]string, 0, len(ordered))
	for _, leaf := range ordered {
		key := leaf.LeafRole + "\x00" + leaf.LeafID
		if _, duplicate := seen[key]; duplicate {
			return nil, fmt.Errorf("duplicate leaf id %s", leaf.LeafID)
		}
		seen[key] = struct{}{}
		if err := assertProvenanceDigest("leaf.commitment", leaf.Commitment); err != nil {
			return nil, err
		}
		digests = append(digests, leaf.Commitment)
	}
	return digests, nil
}

// TopLeafDigest builds a blinded top-tree leaf. The randomizer keeps the top
// tree from leaking which subtrees are empty or shared between capsules.
func TopLeafDigest(leafRole, commitment, randomizer string) (string, error) {
	known := false
	for _, role := range TopLeafRoles {
		if role == leafRole {
			known = true
			break
		}
	}
	if !known {
		return "", fmt.Errorf("unknown top leaf role: %q", leafRole)
	}
	if err := assertProvenanceDigest("commitment", commitment); err != nil {
		return "", err
	}
	if err := assertRandomizer(randomizer); err != nil {
		return "", err
	}
	return DomainHashHex("attesto.provenance.v1.capsule_leaf", map[string]any{
		"leaf_role":  leafRole,
		"commitment": commitment,
		"randomizer": randomizer,
	})
}

// CapsuleRoot folds the six typed top leaves into the capsule root.
func CapsuleRoot(orderedTopLeafDigests []string) (string, error) {
	if len(orderedTopLeafDigests) != len(TopLeafRoles) {
		return "", fmt.Errorf("capsule tree must carry exactly %d leaves", len(TopLeafRoles))
	}
	merkleRoot, err := provenanceFold("attesto.provenance.v1.capsule_node", orderedTopLeafDigests)
	if err != nil {
		return "", err
	}
	return DomainHashHex("attesto.provenance.v1.capsule_root", map[string]any{
		"kind":        "root",
		"leaf_count":  len(orderedTopLeafDigests),
		"merkle_root": merkleRoot,
	})
}

// ProvenanceProofStep is one sibling hop in an inclusion proof.
type ProvenanceProofStep struct {
	Side    string `json:"side"`
	Sibling string `json:"sibling"`
}

// TwoHopProof is a complete disclosure: one subtree leaf, proven to the capsule
// root.
type TwoHopProof struct {
	Subtree           string                `json:"subtree"`
	Leaf              string                `json:"leaf"`
	SubtreeSteps      []ProvenanceProofStep `json:"subtree_steps"`
	SubtreeLeafCount  int                   `json:"subtree_leaf_count"`
	SubtreeRoot       string                `json:"subtree_root"`
	TopLeafRole       string                `json:"top_leaf_role"`
	TopLeafRandomizer string                `json:"top_leaf_randomizer"`
	TopSteps          []ProvenanceProofStep `json:"top_steps"`
	CapsuleRoot       string                `json:"capsule_root"`
}

func replayProvenanceProof(domain, leaf string, steps []ProvenanceProofStep) (string, error) {
	current := leaf
	for _, step := range steps {
		if err := assertProvenanceDigest("proof.sibling", step.Sibling); err != nil {
			return "", err
		}
		var err error
		switch step.Side {
		case "right":
			current, err = provenanceNode(domain, current, step.Sibling)
		case "left":
			current, err = provenanceNode(domain, step.Sibling, current)
		default:
			return "", fmt.Errorf("unknown proof side: %q", step.Side)
		}
		if err != nil {
			return "", err
		}
	}
	return current, nil
}

// VerifyTwoHop verifies a disclosure: leaf -> subtree root -> capsule root.
//
// It returns (false, nil) for a cryptographic failure and an error for a
// malformed object, so a caller can tell "this proof does not hold" from "this
// object is not a proof".
//
// The cross-tree attack this refuses: presenting a claim leaf against
// evidence_root. It fails on two independent grounds — the subtree folds under
// a different node domain, and the top leaf binds leaf_role.
func VerifyTwoHop(proof TwoHopProof) (bool, error) {
	treeDomain, ok := subtreeTreeDomain[proof.Subtree]
	if !ok {
		return false, fmt.Errorf("unknown subtree: %q", proof.Subtree)
	}
	for field, digest := range map[string]string{
		"leaf":         proof.Leaf,
		"subtree_root": proof.SubtreeRoot,
		"capsule_root": proof.CapsuleRoot,
	} {
		if err := assertProvenanceDigest(field, digest); err != nil {
			return false, err
		}
	}

	if proof.TopLeafRole != subtreeTopRole[proof.Subtree] {
		return false, nil
	}

	merkleRoot, err := replayProvenanceProof(treeDomain, proof.Leaf, proof.SubtreeSteps)
	if err != nil {
		return false, err
	}
	derivedSubtreeRoot, err := SubtreeRoot(proof.Subtree, merkleRoot, proof.SubtreeLeafCount)
	if err != nil {
		return false, err
	}
	if derivedSubtreeRoot != proof.SubtreeRoot {
		return false, nil
	}

	leaf, err := TopLeafDigest(proof.TopLeafRole, proof.SubtreeRoot, proof.TopLeafRandomizer)
	if err != nil {
		return false, err
	}
	topMerkle, err := replayProvenanceProof("attesto.provenance.v1.capsule_node", leaf, proof.TopSteps)
	if err != nil {
		return false, err
	}
	derived, err := DomainHashHex("attesto.provenance.v1.capsule_root", map[string]any{
		"kind":        "root",
		"leaf_count":  len(TopLeafRoles),
		"merkle_root": topMerkle,
	})
	if err != nil {
		return false, err
	}
	return subtle.ConstantTimeCompare([]byte(derived), []byte(proof.CapsuleRoot)) == 1, nil
}

// EnvelopeCoreCanonicalBytes returns the canonical bytes of the §8.3 envelope
// core: the egress envelope with signature and boundary removed, because both
// bind it and neither can be part of what they bind.
func EnvelopeCoreCanonicalBytes(envelope map[string]any) ([]byte, error) {
	core := make(map[string]any, len(envelope))
	for key, value := range envelope {
		if key == "signature" || key == "boundary" {
			continue
		}
		core[key] = value
	}
	if err := AssertCommitmentSafeNumbers(core, "$"); err != nil {
		return nil, err
	}
	return CanonicalJSON(core)
}