package attesto

// The bundle provenance tree (ADR-0014, Option C) and the frozen revocation
// rule an offline verifier applies through it.
//
// A verifier bundle over a provenance stream commits to the capsule roots it
// spans through one Merkle root. Everything here is a client of
// edge/src/bundle_tree.rs; the bundle-tree-* vectors pin the agreement. The
// revocation rule mirrors the platform's key_revocation.evaluate exactly, so
// the ingest path and an offline verifier reach the same verdict from the same
// facts.

import (
	"crypto/subtle"
	"encoding/json"
	"fmt"
	"sort"
	"time"
)

const (
	BundleTreeDomain         = "attesto.provenance.v1.bundle_tree"
	bundleTreeName           = "bundle_provenance"
	BundleInclusionKind      = "bundle_provenance_inclusion"
	bundleProvenanceRootKey  = "provenance_root"
	bundleProvenanceCountKey = "provenance_event_count"
	bundleKeyLifecycleKey    = "vault_key_lifecycle"

	KeyStatusValid               = "valid"
	KeyStatusRevokedAtReceipt    = "revoked_at_receipt"
	KeyStatusUnknownInstallation = "unknown_installation"
	KeyStatusNotEvaluated        = "not_evaluated"
	FlagSuspectBackdated         = "suspect_backdated"
	// RevocationReasonUnrecorded marks an instant migration reconstructed rather
	// than measured. The verdict stands; the caller is told not to read the
	// instant as measured.
	RevocationReasonUnrecorded = "unrecorded"

	InclusionValid   = "VALID"
	InclusionInvalid = "INVALID"
)

// BundleLeaf is one provenance event as the bundle tree commits to it. The
// installation, key, assurance and vault-claimed occurred_at are bound with the
// capsule root on purpose: revocation is applied to the installation that
// produced the capsule, and a leaf carrying only the root would let that
// installation be swapped under it.
type BundleLeaf struct {
	SeqNo          int64  `json:"seq_no"`
	CapsuleRoot    string `json:"capsule_root"`
	InstallationID string `json:"installation_id"`
	KeyID          string `json:"key_id"`
	VaultAssurance string `json:"vault_assurance"`
	OccurredAt     string `json:"occurred_at"`
}

// BundleProvenanceTree is the committed tree plus what a prover needs.
type BundleProvenanceTree struct {
	LeafCount          int
	MerkleRoot         string
	ProvenanceRoot     string
	OrderedLeaves      []BundleLeaf
	OrderedLeafDigests []string
}

// BundleInclusionProof is one leaf, proven to the typed provenance root.
type BundleInclusionProof struct {
	Leaf           BundleLeaf            `json:"leaf"`
	LeafCount      int                   `json:"leaf_count"`
	Steps          []ProvenanceProofStep `json:"steps"`
	ProvenanceRoot string                `json:"provenance_root"`
}

func bundleLeafValue(leaf BundleLeaf) (map[string]any, error) {
	if leaf.SeqNo < 0 {
		return nil, fmt.Errorf("bundle leaf seq_no must be a non-negative integer")
	}
	if leaf.InstallationID == "" || leaf.KeyID == "" || leaf.OccurredAt == "" {
		return nil, fmt.Errorf("bundle leaf installation_id, key_id and occurred_at must be non-empty")
	}
	known := false
	for _, level := range VaultAssuranceLevels {
		if level == leaf.VaultAssurance {
			known = true
			break
		}
	}
	if !known {
		// L3 included: it is verifier-derived and has no on-wire form, so a
		// leaf claiming it is malformed rather than merely invalid.
		return nil, fmt.Errorf("bundle leaf vault_assurance must be one of %v", VaultAssuranceLevels)
	}
	if err := assertProvenanceDigest("leaf.capsule_root", leaf.CapsuleRoot); err != nil {
		return nil, err
	}
	return map[string]any{
		"kind":            "leaf",
		"seq_no":          leaf.SeqNo,
		"capsule_root":    leaf.CapsuleRoot,
		"installation_id": leaf.InstallationID,
		"key_id":          leaf.KeyID,
		"vault_assurance": leaf.VaultAssurance,
		"occurred_at":     leaf.OccurredAt,
	}, nil
}

// BundleProvenanceLeaf hashes one provenance event as the bundle tree commits
// to it.
func BundleProvenanceLeaf(leaf BundleLeaf) (string, error) {
	value, err := bundleLeafValue(leaf)
	if err != nil {
		return "", err
	}
	if err := AssertCommitmentSafeNumbers(value, "$"); err != nil {
		return "", err
	}
	return DomainHashHex(BundleTreeDomain, value)
}

func bundleTypedRoot(leafCount int, merkleRoot string) (string, error) {
	return DomainHashHex(BundleTreeDomain, map[string]any{
		"kind":        "root",
		"tree":        bundleTreeName,
		"leaf_count":  leafCount,
		"merkle_root": merkleRoot,
	})
}

// BundleProvenanceRoot folds the events a bundle spans, in seq_no order, into
// its typed root. A repeated seq_no is refused: one event cannot be two leaves.
func BundleProvenanceRoot(leaves []BundleLeaf) (*BundleProvenanceTree, error) {
	if len(leaves) == 0 {
		return nil, fmt.Errorf("cannot build an empty %s tree", bundleTreeName)
	}
	ordered := append([]BundleLeaf(nil), leaves...)
	sort.SliceStable(ordered, func(left, right int) bool {
		return ordered[left].SeqNo < ordered[right].SeqNo
	})
	digests := make([]string, 0, len(ordered))
	for index, leaf := range ordered {
		if index > 0 && ordered[index-1].SeqNo == leaf.SeqNo {
			return nil, fmt.Errorf("duplicate leaf id %d", leaf.SeqNo)
		}
		digest, err := BundleProvenanceLeaf(leaf)
		if err != nil {
			return nil, err
		}
		digests = append(digests, digest)
	}
	merkleRoot, err := provenanceFold(BundleTreeDomain, digests)
	if err != nil {
		return nil, err
	}
	root, err := bundleTypedRoot(len(digests), merkleRoot)
	if err != nil {
		return nil, err
	}
	return &BundleProvenanceTree{
		LeafCount:          len(digests),
		MerkleRoot:         merkleRoot,
		ProvenanceRoot:     root,
		OrderedLeaves:      ordered,
		OrderedLeafDigests: digests,
	}, nil
}

func collectProvenanceProof(domain string, level []string, index int) ([]ProvenanceProofStep, error) {
	steps := []ProvenanceProofStep{}
	current := append([]string(nil), level...)
	for len(current) > 1 {
		next := make([]string, 0, (len(current)+1)/2)
		nextIndex := index
		for cursor := 0; cursor < len(current); cursor += 2 {
			if cursor+1 >= len(current) {
				// Promoted node: it rises with no sibling, so no proof step.
				if cursor == index {
					nextIndex = len(next)
				}
				next = append(next, current[cursor])
				continue
			}
			if cursor == index {
				steps = append(steps, ProvenanceProofStep{Side: "right", Sibling: current[cursor+1]})
				nextIndex = len(next)
			} else if cursor+1 == index {
				steps = append(steps, ProvenanceProofStep{Side: "left", Sibling: current[cursor]})
				nextIndex = len(next)
			}
			node, err := provenanceNode(domain, current[cursor], current[cursor+1])
			if err != nil {
				return nil, err
			}
			next = append(next, node)
		}
		current = next
		index = nextIndex
	}
	return steps, nil
}

// BundleProvenanceProof proves one event under the bundle's provenance root.
func BundleProvenanceProof(leaves []BundleLeaf, seqNo int64) (*BundleInclusionProof, error) {
	tree, err := BundleProvenanceRoot(leaves)
	if err != nil {
		return nil, err
	}
	for index, leaf := range tree.OrderedLeaves {
		if leaf.SeqNo != seqNo {
			continue
		}
		steps, err := collectProvenanceProof(BundleTreeDomain, tree.OrderedLeafDigests, index)
		if err != nil {
			return nil, err
		}
		return &BundleInclusionProof{
			Leaf:           leaf,
			LeafCount:      tree.LeafCount,
			Steps:          steps,
			ProvenanceRoot: tree.ProvenanceRoot,
		}, nil
	}
	return nil, fmt.Errorf("unknown seq_no: %d", seqNo)
}

// VerifyBundleProvenanceInclusion checks that leaf sits under provenanceRoot.
//
// The leaf is re-hashed from its fields, never taken as a digest, so a proof
// cannot substitute one between leaf and root. It returns (false, nil) for a
// cryptographic failure and an error for a malformed object.
func VerifyBundleProvenanceInclusion(provenanceRoot string, leaf BundleLeaf, steps []ProvenanceProofStep, leafCount int) (bool, error) {
	if err := assertProvenanceDigest("provenance_root", provenanceRoot); err != nil {
		return false, err
	}
	if leafCount < 1 {
		return false, fmt.Errorf("leaf_count must be a positive integer")
	}
	digest, err := BundleProvenanceLeaf(leaf)
	if err != nil {
		return false, err
	}
	merkleRoot, err := replayProvenanceProof(BundleTreeDomain, digest, steps)
	if err != nil {
		return false, err
	}
	derived, err := bundleTypedRoot(leafCount, merkleRoot)
	if err != nil {
		return false, err
	}
	return subtle.ConstantTimeCompare([]byte(derived), []byte(provenanceRoot)) == 1, nil
}

// KeyRevocationVerdict is what the key lifecycle says about one event, and why.
type KeyRevocationVerdict struct {
	Status    string     `json:"status"`
	Flags     []string   `json:"flags"`
	RevokedAt *time.Time `json:"revoked_at"`
	Reason    string     `json:"reason"`
	// True when the effective instant was reconstructed by migration. The
	// verdict still stands; the caller is told not to read it as measured.
	InstantReconstructed bool `json:"instant_reconstructed"`
}

// Accepted reports whether the key was live at receipt. It says nothing about
// the signature, which is checked separately.
func (v KeyRevocationVerdict) Accepted() bool { return v.Status == KeyStatusValid }

// EvaluateKeyRevocation decides whether a key was live when the platform
// received the event.
//
// Revocation is evaluated against the platform receipt time, never the
// vault-claimed occurred_at: a holder controls what it claims, not when the
// platform received it. The boundary is inclusive — an event receipted exactly
// at the revocation instant is revoked, because the alternative gives a
// compromised key one more accepted event. A claim that predates revocation
// while its receipt does not is flagged suspect_backdated in addition to being
// revoked, not instead of.
//
// A nil revokedAt means the key was never revoked, which is a different thing
// from a key revoked in the future and must not be conflated: the second is a
// scheduled retirement and is still evidence.
func EvaluateKeyRevocation(revokedAt *time.Time, receiptTime time.Time, claimedOccurredAt *time.Time, reason string) KeyRevocationVerdict {
	if revokedAt == nil {
		return KeyRevocationVerdict{Status: KeyStatusValid, Flags: []string{}}
	}
	effective := revokedAt.UTC()
	received := receiptTime.UTC()
	reconstructed := reason == RevocationReasonUnrecorded
	if received.Before(effective) {
		return KeyRevocationVerdict{
			Status:               KeyStatusValid,
			Flags:                []string{},
			RevokedAt:            &effective,
			Reason:               reason,
			InstantReconstructed: reconstructed,
		}
	}
	flags := []string{KeyStatusRevokedAtReceipt}
	if claimedOccurredAt != nil && claimedOccurredAt.UTC().Before(effective) {
		flags = append(flags, FlagSuspectBackdated)
	}
	return KeyRevocationVerdict{
		Status:               KeyStatusRevokedAtReceipt,
		Flags:                flags,
		RevokedAt:            &effective,
		Reason:               reason,
		InstantReconstructed: reconstructed,
	}
}

// BundleProvenanceNotClaimed states what a verified inclusion does not prove.
var BundleProvenanceNotClaimed = []map[string]string{
	{
		"id": "bundle_asserts_capsule_existence_not_contents",
		"statement": "The provenance_root proves this capsule root was among the events " +
			"the bundle spans. It says nothing about what the capsule contains; the " +
			"platform never opens one.",
	},
	{
		"id": "revocation_evaluated_against_platform_receipt_time_not_vault_occurred_at",
		"statement": "Key revocation is evaluated against the platform receipt time of " +
			"this seq_no. The vault-claimed occurred_at is reported, never trusted to " +
			"escape revocation.",
	},
	{
		"id": "key_lifecycle_as_of_bundle_build",
		"statement": "vault_key_lifecycle is the installation's lifecycle when the bundle " +
			"was built. A revocation recorded later is not in this bundle.",
	},
}

// BundleProvenanceReport is what one inclusion established against its bundle,
// fact by fact. Inclusion and KeyStatus are separate on purpose: a capsule root
// can be provably under the bundle while its key was revoked before receipt,
// and collapsing the two would lose exactly the distinction an investigator
// needs. Ok is true only when the bundle hash holds, the inclusion verifies and
// the key was live at receipt.
type BundleProvenanceReport struct {
	Ok             bool                `json:"ok"`
	Inclusion      string              `json:"inclusion"`
	KeyStatus      string              `json:"key_status"`
	Flags          []string            `json:"flags"`
	SeqNo          *int64              `json:"seq_no"`
	InstallationID string              `json:"installation_id"`
	CapsuleRoot    string              `json:"capsule_root"`
	ReceiptTime    string              `json:"receipt_time"`
	RevokedAt      string              `json:"revoked_at"`
	Problems       []string            `json:"problems"`
	NotClaimed     []map[string]string `json:"not_claimed"`
}

func parseRFC3339(raw string) (*time.Time, bool) {
	parsed, err := time.Parse(time.RFC3339Nano, raw)
	if err != nil {
		return nil, false
	}
	return &parsed, true
}

func receiptIssuedAt(bundle map[string]any, seqNo int64) string {
	for _, raw := range asSlice(bundle["receipts"]) {
		item, ok := raw.(map[string]any)
		if !ok {
			continue
		}
		itemSeq, ok := item["seq_no"].(float64)
		if !ok || int64(itemSeq) != seqNo {
			continue
		}
		receipt, _ := item["receipt"].(map[string]any)
		payload, _ := receipt["payload"].(map[string]any)
		return toString(payload["issued_at"])
	}
	return ""
}

// VerifyBundleProvenance verifies one capsule's inclusion in a verifier bundle,
// offline. It needs nothing but the bundle and the inclusion object: it
// recomputes the bundle hash so the provenance root and key lifecycle are
// authenticated, checks the leaf under the root, takes the receipt time for the
// leaf's seq_no from the bundle's own receipts, and applies the frozen
// revocation rule to the installation the leaf names. Every problem is
// collected rather than returned on the first one.
func VerifyBundleProvenance(bundle map[string]any, inclusion map[string]any) BundleProvenanceReport {
	problems := []string{}
	payload, ok := bundle["payload"].(map[string]any)
	if !ok {
		payload = map[string]any{}
		problems = append(problems, "invalid bundle object")
	}
	bundleHash := toString(bundle["bundle_hash"])
	if len(payload) > 0 {
		derived, err := DomainHashHex(ProofstreamDomains["bundle"], payload)
		if err != nil || derived != bundleHash {
			problems = append(problems, "bundle_hash mismatch")
		}
	}

	var leaf BundleLeaf
	var seqNo *int64
	if rawLeaf, ok := inclusion["leaf"].(map[string]any); ok {
		if encoded, err := json.Marshal(rawLeaf); err == nil {
			_ = json.Unmarshal(encoded, &leaf)
		}
		if value, ok := rawLeaf["seq_no"].(float64); ok {
			n := int64(value)
			seqNo = &n
		}
	}

	if toString(inclusion["kind"]) != BundleInclusionKind {
		problems = append(problems, "inclusion is not a bundle_provenance_inclusion object")
	}
	if toString(inclusion["bundle_hash"]) != bundleHash {
		problems = append(problems, "inclusion is for a different bundle")
	}

	included := InclusionInvalid
	declaredRoot, hasRoot := payload[bundleProvenanceRootKey].(string)
	if !hasRoot {
		problems = append(problems, "bundle carries no provenance_root")
	} else {
		if toString(inclusion["provenance_root"]) != declaredRoot {
			problems = append(problems, "inclusion names a different provenance_root")
		}
		leafCount, _ := inclusion["leaf_count"].(float64)
		declaredCount, _ := payload[bundleProvenanceCountKey].(float64)
		if leafCount != declaredCount {
			problems = append(problems, "inclusion leaf_count does not match provenance_event_count")
		}
		var steps []ProvenanceProofStep
		if encoded, err := json.Marshal(inclusion["steps"]); err == nil {
			_ = json.Unmarshal(encoded, &steps)
		}
		verified, err := VerifyBundleProvenanceInclusion(declaredRoot, leaf, steps, int(leafCount))
		switch {
		case err != nil:
			problems = append(problems, "inclusion is malformed: "+err.Error())
		case verified:
			included = InclusionValid
		default:
			problems = append(problems, "leaf is not included under the bundle's provenance_root")
		}
	}

	receiptTime := ""
	if seqNo != nil {
		receiptTime = receiptIssuedAt(bundle, *seqNo)
	}
	if receiptTime == "" {
		problems = append(problems, fmt.Sprintf("bundle carries no receipt for seq_no %v", formatSeqNo(seqNo)))
	}

	var entry map[string]any
	for _, raw := range asSlice(payload[bundleKeyLifecycleKey]) {
		candidate, ok := raw.(map[string]any)
		if ok && toString(candidate["installation_id"]) == leaf.InstallationID && leaf.InstallationID != "" {
			entry = candidate
			break
		}
	}

	keyStatus := KeyStatusUnknownInstallation
	flags := []string{}
	revokedAt := ""
	switch {
	case entry == nil:
		problems = append(problems, fmt.Sprintf("installation %s is not in the bundle's key lifecycle", leaf.InstallationID))
	case receiptTime == "":
		keyStatus = KeyStatusNotEvaluated
	default:
		if toString(entry["key_id"]) != leaf.KeyID {
			problems = append(problems, "key lifecycle key_id does not match the leaf")
		}
		received, ok := parseRFC3339(receiptTime)
		if !ok {
			keyStatus = KeyStatusNotEvaluated
			problems = append(problems, "receipt issued_at is not an RFC 3339 instant")
			break
		}
		var effective *time.Time
		if raw := toString(entry["revoked_at"]); raw != "" {
			parsed, ok := parseRFC3339(raw)
			if !ok {
				keyStatus = KeyStatusNotEvaluated
				problems = append(problems, "key lifecycle is malformed: revoked_at is not an RFC 3339 instant")
				break
			}
			effective = parsed
			revokedAt = raw
		}
		claimed, _ := parseRFC3339(leaf.OccurredAt)
		verdict := EvaluateKeyRevocation(effective, *received, claimed, toString(entry["revocation_reason"]))
		keyStatus = verdict.Status
		flags = verdict.Flags
		if !verdict.Accepted() {
			problems = append(problems, fmt.Sprintf(
				"installation %s was revoked before seq_no %s was received", leaf.InstallationID, formatSeqNo(seqNo),
			))
		}
	}

	return BundleProvenanceReport{
		Ok:             len(problems) == 0 && included == InclusionValid && keyStatus == KeyStatusValid,
		Inclusion:      included,
		KeyStatus:      keyStatus,
		Flags:          flags,
		SeqNo:          seqNo,
		InstallationID: leaf.InstallationID,
		CapsuleRoot:    leaf.CapsuleRoot,
		ReceiptTime:    receiptTime,
		RevokedAt:      revokedAt,
		Problems:       problems,
		NotClaimed:     BundleProvenanceNotClaimed,
	}
}

func formatSeqNo(seqNo *int64) string {
	if seqNo == nil {
		return "<none>"
	}
	return fmt.Sprintf("%d", *seqNo)
}