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"
	"strings"
)

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":     {},
	"attesto.disclosure.v2":                 {},
	"attesto.zk.range.v1.statement":         {},
	"attesto.zk.range.v1.transcript":        {},
}

// 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,
	})
}

// OrderedTopLeafDigests builds the six typed top leaves, requiring each role
// exactly once.
//
// CapsuleRoot receives digests, so by then a role is no longer visible and a
// tree carrying evidence_root twice with vault_identity_commitment missing folds
// to a root it will accept. The check has to happen here, where the roles still
// exist, which is also why callers should reach for this rather than assembling
// the slice themselves.
func OrderedTopLeafDigests(commitments, randomizers map[string]string) ([]string, error) {
	for _, role := range TopLeafRoles {
		if _, ok := commitments[role]; !ok {
			return nil, fmt.Errorf("capsule tree is missing top leaf %s", role)
		}
	}
	known := make(map[string]struct{}, len(TopLeafRoles))
	for _, role := range TopLeafRoles {
		known[role] = struct{}{}
	}
	for role := range commitments {
		if _, ok := known[role]; !ok {
			return nil, fmt.Errorf("unknown top leaf role: %s", role)
		}
	}
	digests := make([]string, 0, len(TopLeafRoles))
	for _, role := range TopLeafRoles {
		digest, err := TopLeafDigest(role, commitments[role], randomizers[role])
		if err != nil {
			return nil, err
		}
		digests = append(digests, digest)
	}
	return digests, nil
}

// 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)
}

// ---------------------------------------------------------------- predicates

const (
	ZKRangeProtocol              = "ATTESTO-ZK-RANGE-001"
	ZKRangeProtocolVersion       = "0.1"
	PredicateResultSchema        = "attesto.provenance.predicate_result"
	PredicateResultSchemaVersion = "0.1"
)

// RequiredNonClaims must appear in every predicate result. A result that dropped
// one would be read as the stronger statement, which is the failure this
// vocabulary prevents.
var RequiredNonClaims = [3]string{
	"detector_correctness_not_proven",
	"content_truth_not_proven",
	"ai_generation_not_proven",
}

// forbiddenResultFields would let a consumer render a proven bound as a verdict
// about the content. A range proof says a named detector's measurement fell
// inside an interval and nothing more.
var forbiddenResultFields = map[string]struct{}{
	"ai_generated": {}, "synthetic": {}, "is_fake": {}, "authentic": {},
	"confidence": {}, "score": {}, "probability": {},
}

// PredicateReport separates what this client checked from what the issuer claims.
type PredicateReport struct {
	Protocol         string            `json:"protocol"`
	CapsuleRoot      string            `json:"capsule_root"`
	ClaimID          string            `json:"claim_id"`
	CommitmentC      string            `json:"commitment_c"`
	Predicate        map[string]any    `json:"predicate"`
	VerifiedHere     map[string]string `json:"verified_here"`
	ReportedByIssuer map[string]any    `json:"reported_by_issuer"`
	NotClaimed       []map[string]any  `json:"not_claimed"`
}

func rejectVerdictFields(node any, path string) error {
	switch typed := node.(type) {
	case map[string]any:
		for key, value := range typed {
			if _, bad := forbiddenResultFields[strings.ToLower(key)]; bad {
				return fmt.Errorf("predicate result carries a verdict-shaped field at %s.%s", path, key)
			}
			if err := rejectVerdictFields(value, path+"."+key); err != nil {
				return err
			}
		}
	case []any:
		for index, value := range typed {
			if err := rejectVerdictFields(value, fmt.Sprintf("%s[%d]", path, index)); err != nil {
				return err
			}
		}
	}
	return nil
}

// InspectPredicateResult reports what this SDK established, kept apart from what
// the issuer claims.
//
// This is a verification client without ristretto255 arithmetic, so it cannot
// check a range proof. It says not_checked rather than passing the issuer's word
// through as though it had verified it: an SDK that reported the issuer's
// "verified" as its own is the failure this construction exists to prevent.
//
// capsuleInclusion is nil when the caller did not check inclusion.
func InspectPredicateResult(result map[string]any, capsuleInclusion *bool) (*PredicateReport, error) {
	if result["schema"] != PredicateResultSchema || result["schema_version"] != PredicateResultSchemaVersion {
		return nil, fmt.Errorf("unsupported predicate result schema")
	}
	if result["protocol"] != ZKRangeProtocol || result["protocol_version"] != ZKRangeProtocolVersion {
		return nil, fmt.Errorf("unsupported predicate protocol")
	}

	rawClaims, _ := result["not_claimed"].([]any)
	declared := map[string]struct{}{}
	notClaimed := make([]map[string]any, 0, len(rawClaims))
	for _, raw := range rawClaims {
		claim, ok := raw.(map[string]any)
		if !ok {
			return nil, fmt.Errorf("predicate result carries a malformed non-claim")
		}
		if id, ok := claim["id"].(string); ok {
			declared[id] = struct{}{}
		}
		copied := map[string]any{}
		for key, value := range claim {
			copied[key] = value
		}
		notClaimed = append(notClaimed, copied)
	}
	for _, required := range RequiredNonClaims {
		if _, ok := declared[required]; !ok {
			return nil, fmt.Errorf("predicate result omits required non-claims: %s", required)
		}
	}

	if err := rejectVerdictFields(result, "result"); err != nil {
		return nil, err
	}

	bound := map[string]string{}
	for _, field := range []string{"capsule_root", "claim_id", "commitment_c"} {
		value, ok := result[field].(string)
		if !ok || value == "" {
			return nil, fmt.Errorf("predicate result has no %s to bind to", field)
		}
		bound[field] = value
	}

	inclusion := "not_checked"
	if capsuleInclusion != nil {
		if *capsuleInclusion {
			inclusion = "verified"
		} else {
			inclusion = "failed"
		}
	}

	predicate, _ := result["predicate"].(map[string]any)
	issuer, _ := result["verification"].(map[string]any)
	return &PredicateReport{
		Protocol:    ZKRangeProtocol,
		CapsuleRoot: bound["capsule_root"],
		ClaimID:     bound["claim_id"],
		CommitmentC: bound["commitment_c"],
		Predicate:   predicate,
		// zk_predicate is always not_checked: verifying the proof needs curve
		// arithmetic this client does not carry.
		VerifiedHere: map[string]string{
			"zk_predicate":      "not_checked",
			"capsule_inclusion": inclusion,
		},
		// What the issuer says it checked, kept separate so a reader can tell a
		// claim from a check.
		ReportedByIssuer: issuer,
		NotClaimed:       notClaimed,
	}, nil
}

// ZKRangeWidths are the proof widths ATTESTO-ZK-RANGE-001 permits. Pinned rather
// than derived: a client offering a width outside this set would build
// statements the proving library refuses after the transcript is already bound.
var ZKRangeWidths = [4]uint{8, 16, 32, 64}

// zkRangeStatementFields is exactly what a range statement carries. Every field
// is folded into the proof transcript, so an extra one would bind to nothing and
// a missing one would change the challenges.
var zkRangeStatementFields = map[string]struct{}{
	"capsule_root": {}, "claim_id": {}, "claim_descriptor_version": {},
	"provider_id": {}, "provider_version": {}, "commitment_c": {},
	"predicate_type": {}, "lower_bound": {}, "upper_bound": {},
	"encoding_version": {}, "proof_scheme_version": {}, "verifier_nonce": {},
}

// ZKRangeWidth returns the smallest permitted width covering the whole interval.
//
// Both proved differences are bounded by upper-lower, so one width serves both.
// It is derived from the public bounds and never chosen by the prover: a prover
// who picked it could prove a wider range than the statement says.
func ZKRangeWidth(lowerBound, upperBound uint64) (uint, error) {
	if upperBound < lowerBound {
		return 0, fmt.Errorf("upper bound is below its lower bound")
	}
	span := upperBound - lowerBound
	for _, width := range ZKRangeWidths {
		if width == 64 || span < (uint64(1)<<width) {
			return width, nil
		}
	}
	return 0, fmt.Errorf("interval is wider than the largest permitted proof")
}

// encodedBound reads a JSON number as the encoded integer it must be. A float
// bound is refused rather than truncated: the encoding registry exists so no
// rounding step is left for three SDKs to disagree about.
func encodedBound(value any, name string) (uint64, error) {
	number, ok := value.(float64)
	if !ok {
		return 0, fmt.Errorf("%s must be an encoded integer, never a float", name)
	}
	if number < 0 || number != float64(uint64(number)) {
		return 0, fmt.Errorf("%s must be an encoded integer, never a float", name)
	}
	return uint64(number), nil
}

// ValidateRangeStatement checks a statement is well-formed and returns the width
// its bounds imply.
//
// This does not verify the proof — that needs curve arithmetic no SDK carries.
// It refuses the statements no honest prover produced, which is a check a
// verification client can make on its own.
func ValidateRangeStatement(statement map[string]any) (uint, error) {
	for field := range zkRangeStatementFields {
		if _, ok := statement[field]; !ok {
			return 0, fmt.Errorf("range statement is missing %s", field)
		}
	}
	for field := range statement {
		if _, ok := zkRangeStatementFields[field]; !ok {
			return 0, fmt.Errorf("range statement carries an unexpected field: %s", field)
		}
	}
	if statement["predicate_type"] != "inclusive_range" {
		return 0, fmt.Errorf("statement predicate is not the v1 inclusive range")
	}
	for _, field := range []string{
		"capsule_root", "claim_id", "claim_descriptor_version", "provider_id",
		"provider_version", "commitment_c", "encoding_version",
		"proof_scheme_version", "verifier_nonce",
	} {
		value, ok := statement[field].(string)
		if !ok || value == "" {
			return 0, fmt.Errorf("range statement %s is empty or not a string", field)
		}
	}
	lower, err := encodedBound(statement["lower_bound"], "lower_bound")
	if err != nil {
		return 0, err
	}
	upper, err := encodedBound(statement["upper_bound"], "upper_bound")
	if err != nil {
		return 0, err
	}
	return ZKRangeWidth(lower, upper)
}