proofborne-core 0.2.0-alpha.1

use std::collections::BTreeSet;

use serde::{Deserialize, Serialize};
use thiserror::Error;

use crate::{ContextSource, MemoryStatus, SCHEMA_VERSION, hash_json};

/// Version of the deterministic context-compaction algorithm.
pub const COMPACTION_ALGORITHM_VERSION: &str = "proofborne.compaction.v1";
/// Deterministic fallback estimator used when an adapter exposes no model tokenizer.
pub const COMPACTION_TOKEN_ESTIMATOR: &str = "utf8_ascii_div_3_plus_non_ascii_scalars.v1";

/// Hard provider-context boundary used by one compaction decision.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct CompactionBudget {
    /// Provider-declared context-window size.
    pub context_tokens: u64,
    /// Tokens reserved for the provider response.
    pub reserved_output_tokens: u64,
    /// Maximum deterministic fallback token estimate available to the input prompt.
    pub max_input_tokens: u64,
    /// Independent serialized-byte ceiling for the input prompt.
    pub max_input_bytes: u64,
}

impl CompactionBudget {
    /// Creates a budget whose input allowance is the provider window minus the response reserve.
    pub fn new(
        context_tokens: u64,
        reserved_output_tokens: u64,
        max_input_bytes: u64,
    ) -> Result<Self, CompactionError> {
        let max_input_tokens = context_tokens
            .checked_sub(reserved_output_tokens)
            .filter(|available| *available > 0)
            .ok_or(CompactionError::InvalidBudget)?;
        let budget = Self {
            context_tokens,
            reserved_output_tokens,
            max_input_tokens,
            max_input_bytes,
        };
        budget.validate()?;
        Ok(budget)
    }

    /// Validates the arithmetic and byte boundary committed by this budget.
    pub fn validate(&self) -> Result<(), CompactionError> {
        if self.context_tokens == 0
            || self.reserved_output_tokens == 0
            || self.max_input_tokens == 0
            || self.max_input_bytes == 0
            || self
                .reserved_output_tokens
                .checked_add(self.max_input_tokens)
                != Some(self.context_tokens)
        {
            return Err(CompactionError::InvalidBudget);
        }
        Ok(())
    }
}

/// Prompt-category accounting measured by the versioned deterministic fallback estimator.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct CompactionUsage {
    /// Versioned estimator used for every token count in this record.
    pub token_estimator: String,
    /// Stable runtime instructions.
    pub instruction_tokens: u64,
    /// Serialized tool definitions.
    pub tool_tokens: u64,
    /// Task contract, claim scope, criteria, waivers, and constraints.
    pub contract_tokens: u64,
    /// Proof obligations, evidence, supersession, and counterevidence.
    pub evidence_tokens: u64,
    /// Provenance-bound memory retrieved for this turn.
    pub memory_tokens: u64,
    /// Provider-neutral recent conversation history.
    pub history_tokens: u64,
    /// Sum of every input category above.
    pub input_tokens: u64,
    /// Exact serialized size of the resulting provider input.
    pub input_bytes: u64,
}
impl Default for CompactionUsage {
    fn default() -> Self {
        Self {
            token_estimator: COMPACTION_TOKEN_ESTIMATOR.to_owned(),
            instruction_tokens: 0,
            tool_tokens: 0,
            contract_tokens: 0,
            evidence_tokens: 0,
            memory_tokens: 0,
            history_tokens: 0,
            input_tokens: 0,
            input_bytes: 0,
        }
    }
}

impl CompactionUsage {
    /// Recomputes the category sum without saturating arithmetic.
    pub fn category_total(&self) -> Option<u64> {
        self.instruction_tokens
            .checked_add(self.tool_tokens)?
            .checked_add(self.contract_tokens)?
            .checked_add(self.evidence_tokens)?
            .checked_add(self.memory_tokens)?
            .checked_add(self.history_tokens)
    }

    /// Validates category accounting against a hard budget.
    pub fn validate(&self, budget: &CompactionBudget) -> Result<(), CompactionError> {
        if self.token_estimator != COMPACTION_TOKEN_ESTIMATOR {
            return Err(CompactionError::UnsupportedTokenEstimator(
                self.token_estimator.clone(),
            ));
        }
        budget.validate()?;
        if self.category_total() != Some(self.input_tokens) {
            return Err(CompactionError::UsageAccounting);
        }
        if self.input_tokens > budget.max_input_tokens || self.input_bytes > budget.max_input_bytes
        {
            return Err(CompactionError::BudgetExceeded);
        }
        Ok(())
    }
}

/// Semantic class of one item considered by compaction.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
#[serde(rename_all = "snake_case")]
pub enum CompactionContextKind {
    /// Complete task contract and claim scope.
    Contract,
    /// One acceptance criterion or open proof obligation.
    Criterion,
    /// Runtime evidence or counterevidence.
    Evidence,
    /// Provenance-bound durable memory.
    Memory,
    /// Provider-neutral conversation item.
    History,
    /// Current workspace generation and state binding.
    Workspace,
    /// Side-effect watermark or ambiguous in-flight action.
    SideEffect,
    /// Crash-recovery continuation state.
    Recovery,
    /// Policy, tool, or producer authority binding.
    Authority,
    /// Public sensitive-data handling metadata; never secret content.
    SecretTaint,
}

/// Deterministic priority class. Lower classes are selected first.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
#[serde(rename_all = "snake_case")]
pub enum CompactionPriority {
    /// Cannot be omitted; overflow blocks the run.
    Pinned,
    /// Active proof material needed to evaluate completion.
    Proof,
    /// Current, provenance-bound memory relevant to the task.
    Retrieved,
    /// Recent provider conversation retained with remaining budget.
    Recent,
}

/// Final disposition of one considered context item.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum CompactionDecision {
    /// Included in the resulting provider context.
    Retained,
    /// Deliberately excluded from the resulting provider context.
    Omitted,
}

/// Runtime-owned reason for a deterministic selection decision.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
#[serde(rename_all = "snake_case")]
pub enum CompactionSelectionReason {
    /// The complete task contract is always pinned.
    TaskContract,
    /// The maximum permitted claim remains explicit.
    ClaimScope,
    /// A required criterion remains explicit.
    RequiredCriterion,
    /// A pending or failed criterion is an open proof obligation.
    OpenObligation,
    /// Failed active evidence still blocks completion.
    UnresolvedCounterevidence,
    /// Successful active evidence is needed to derive criterion state.
    ActiveEvidence,
    /// Human waiver metadata remains reviewable.
    Waiver,
    /// Evidence replacement relationships remain reviewable.
    Supersession,
    /// Current workspace generation and digest remain pinned.
    WorkspaceBinding,
    /// Completed or ambiguous side effects remain pinned.
    SideEffectBinding,
    /// Crash-resume continuation state remains pinned.
    RecoveryBinding,
    /// Policy and tool authority remain pinned.
    AuthorityBinding,
    /// Redaction state is retained without retaining secret contents.
    SecretTaint,
    /// Current memory matched the task contract.
    TaskScopeMatch,
    /// Current memory matched a relevant workspace path.
    PathScopeMatch,
    /// Current memory matched a relevant language symbol.
    SymbolScopeMatch,
    /// Current memory matched a relevant producer.
    ProducerMatch,
    /// Current memory matched deterministic full-text retrieval.
    RetrievalMatch,
    /// A recent conversation item fit the remaining budget.
    RecentHistory,
    /// A non-required item did not fit the remaining budget.
    BudgetExhausted,
    /// Historical memory is excluded unless explicitly requested.
    HistoricalMemoryExcluded,
    /// Evidence or memory was replaced by a newer valid observation.
    Superseded,
}

/// One content-addressed, provenance-preserving selection decision.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct ContextSelection {
    /// Stable namespaced identifier, such as `evidence:<uuid>` or `history:<digest>`.
    pub id: String,
    /// Semantic context class.
    pub kind: CompactionContextKind,
    /// BLAKE3 digest of the exact public serialized item considered.
    pub content_hash: String,
    /// Final inclusion decision.
    pub decision: CompactionDecision,
    /// Runtime-owned deterministic decision reason.
    pub reason: CompactionSelectionReason,
    /// Selection priority.
    pub priority: CompactionPriority,
    /// Whether omission must fail closed instead of degrading context.
    pub required: bool,
    /// Stable total order within the priority class.
    pub rank: u64,
    /// Versioned deterministic fallback token estimate for this item.
    pub estimated_tokens: u64,
    /// Exact serialized public byte count for this item.
    pub byte_count: u64,
    /// Original session/event/evidence/workspace provenance when available.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub source: Option<ContextSource>,
    /// Durable lifecycle for memory selections.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub memory_status: Option<MemoryStatus>,
}

impl ContextSelection {
    fn validate(&self) -> Result<(), CompactionError> {
        if self.id.trim().is_empty() {
            return Err(CompactionError::EmptySelectionId);
        }
        if !valid_digest(&self.content_hash) {
            return Err(CompactionError::InvalidDigest(self.content_hash.clone()));
        }
        if self.estimated_tokens == 0 || self.byte_count == 0 {
            return Err(CompactionError::EmptySelection(self.id.clone()));
        }
        if self.required && self.decision != CompactionDecision::Retained {
            return Err(CompactionError::RequiredSelectionOmitted(self.id.clone()));
        }
        if self.kind == CompactionContextKind::Memory && self.memory_status.is_none() {
            return Err(CompactionError::MemoryLifecycleMissing(self.id.clone()));
        }
        if self.kind != CompactionContextKind::Memory && self.memory_status.is_some() {
            return Err(CompactionError::UnexpectedMemoryLifecycle(self.id.clone()));
        }
        Ok(())
    }
}

/// Complete deterministic decision before provider context is rewritten.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct CompactionPlan {
    /// Public schema identifier.
    pub schema_version: String,
    /// Deterministic algorithm identifier.
    pub algorithm_version: String,
    /// BLAKE3 digest of the uncompacted provider input.
    pub input_hash: String,
    /// BLAKE3 digest of the complete task contract.
    pub contract_hash: String,
    /// Workspace generation to which this decision is bound.
    pub workspace_generation: u64,
    /// Canonical workspace digest at this generation.
    pub state_binding: String,
    /// Runtime side-effect watermark at the decision point.
    pub side_effect_watermark: u64,
    /// Secret-free policy/tool authority digest when configured.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub authority_hash: Option<String>,
    /// Provider-context boundary.
    pub budget: CompactionBudget,
    /// Every considered item in deterministic priority/rank/identifier order.
    pub selections: Vec<ContextSelection>,
}

impl CompactionPlan {
    /// Returns the canonical BLAKE3 digest committed by a receipt.
    pub fn digest(&self) -> Result<String, CompactionError> {
        let value = serde_json::to_value(self).map_err(|_| CompactionError::Serialization)?;
        Ok(hash_json(&value))
    }

    /// Validates schema, bindings, ordering, uniqueness, and fail-closed pinning.
    pub fn validate(&self) -> Result<(), CompactionError> {
        if self.schema_version != SCHEMA_VERSION {
            return Err(CompactionError::UnsupportedSchema(
                self.schema_version.clone(),
            ));
        }
        if self.algorithm_version != COMPACTION_ALGORITHM_VERSION {
            return Err(CompactionError::UnsupportedAlgorithm(
                self.algorithm_version.clone(),
            ));
        }
        for digest in [&self.input_hash, &self.contract_hash, &self.state_binding] {
            if !valid_digest(digest) {
                return Err(CompactionError::InvalidDigest((*digest).clone()));
            }
        }
        if self
            .authority_hash
            .as_deref()
            .is_some_and(|digest| !valid_digest(digest))
        {
            return Err(CompactionError::InvalidDigest(
                self.authority_hash.clone().unwrap_or_default(),
            ));
        }
        self.budget.validate()?;
        if self.selections.is_empty() {
            return Err(CompactionError::EmptyPlan);
        }

        let mut identifiers = BTreeSet::new();
        let mut previous = None;
        for selection in &self.selections {
            selection.validate()?;
            if !identifiers.insert(selection.id.as_str()) {
                return Err(CompactionError::DuplicateSelection(selection.id.clone()));
            }
            let key = (selection.priority, selection.rank, selection.id.as_str());
            if previous.is_some_and(|previous| previous > key) {
                return Err(CompactionError::SelectionOrder);
            }
            previous = Some(key);
        }
        Ok(())
    }
}

/// Durable proof that one validated plan produced one exact provider input.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct CompactionReceipt {
    /// Public schema identifier.
    pub schema_version: String,
    /// Deterministic algorithm identifier.
    pub algorithm_version: String,
    /// Canonical digest of `plan`.
    pub plan_hash: String,
    /// Digest repeated from the plan to make event inspection self-contained.
    pub input_hash: String,
    /// BLAKE3 digest of the exact compacted provider input.
    pub output_hash: String,
    /// BLAKE3 digest of the exact compacted conversation history persisted for recovery.
    pub output_history_hash: String,
    /// Number of conversation items covered by `outputHistoryHash`.
    pub output_history_items: u64,
    /// Workspace generation repeated from the plan.
    pub workspace_generation: u64,
    /// Canonical workspace digest repeated from the plan.
    pub state_binding: String,
    /// Complete decision whose canonical digest is `planHash`.
    pub plan: CompactionPlan,
    /// Exact resulting input accounting.
    pub usage: CompactionUsage,
    /// Sorted identifiers included in the compacted provider context.
    pub retained_ids: Vec<String>,
    /// Sorted identifiers deliberately excluded from the provider context.
    pub omitted_ids: Vec<String>,
}

impl CompactionReceipt {
    /// Validates the plan commitment, output binding, accounting, and identifier sets.
    pub fn validate(&self) -> Result<(), CompactionError> {
        if self.schema_version != SCHEMA_VERSION {
            return Err(CompactionError::UnsupportedSchema(
                self.schema_version.clone(),
            ));
        }
        if self.algorithm_version != COMPACTION_ALGORITHM_VERSION {
            return Err(CompactionError::UnsupportedAlgorithm(
                self.algorithm_version.clone(),
            ));
        }
        self.plan.validate()?;
        if self.plan_hash != self.plan.digest()? {
            return Err(CompactionError::PlanDigest);
        }
        if self.input_hash != self.plan.input_hash
            || self.workspace_generation != self.plan.workspace_generation
            || self.state_binding != self.plan.state_binding
        {
            return Err(CompactionError::PlanBinding);
        }
        for digest in [&self.output_hash, &self.output_history_hash] {
            if !valid_digest(digest) {
                return Err(CompactionError::InvalidDigest((*digest).clone()));
            }
        }
        if self.output_history_items == 0 {
            return Err(CompactionError::PlanBinding);
        }
        self.usage.validate(&self.plan.budget)?;

        let mut retained = self
            .plan
            .selections
            .iter()
            .filter(|selection| selection.decision == CompactionDecision::Retained)
            .map(|selection| selection.id.clone())
            .collect::<Vec<_>>();
        let mut omitted = self
            .plan
            .selections
            .iter()
            .filter(|selection| selection.decision == CompactionDecision::Omitted)
            .map(|selection| selection.id.clone())
            .collect::<Vec<_>>();
        retained.sort();
        omitted.sort();
        if retained != self.retained_ids || omitted != self.omitted_ids {
            return Err(CompactionError::SelectionSets);
        }
        Ok(())
    }
}

fn valid_digest(value: &str) -> bool {
    value.len() == 64
        && value
            .bytes()
            .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
}

/// Compaction contract or receipt validation failure.
#[derive(Debug, Clone, Error, PartialEq, Eq)]
pub enum CompactionError {
    /// Unsupported public schema identifier.
    #[error("unsupported compaction schema: {0}")]
    UnsupportedSchema(String),
    /// Unsupported deterministic algorithm identifier.
    #[error("unsupported compaction algorithm: {0}")]
    UnsupportedAlgorithm(String),
    /// Context budget arithmetic or a hard limit is invalid.
    #[error("compaction budget is invalid")]
    InvalidBudget,
    /// Unsupported token-estimator identifier.
    #[error("unsupported compaction token estimator: {0}")]
    UnsupportedTokenEstimator(String),
    /// Usage categories do not sum to the committed input total.
    #[error("compaction usage accounting is inconsistent")]
    UsageAccounting,
    /// Resulting context exceeds the provider token or byte boundary.
    #[error("compacted context exceeds its hard budget")]
    BudgetExceeded,
    /// A public compaction contract could not be serialized canonically.
    #[error("compaction contract serialization failed")]
    Serialization,
    /// A required digest is not canonical lowercase BLAKE3 hex.
    #[error("invalid compaction digest: {0}")]
    InvalidDigest(String),
    /// A plan contains no considered context.
    #[error("compaction plan must contain at least one selection")]
    EmptyPlan,
    /// A selection identifier is empty.
    #[error("compaction selection identifier must not be empty")]
    EmptySelectionId,
    /// A selection has no measurable public content.
    #[error("compaction selection has empty content: {0}")]
    EmptySelection(String),
    /// A selection identifier occurs more than once.
    #[error("duplicate compaction selection: {0}")]
    DuplicateSelection(String),
    /// Selections are not in deterministic priority/rank/identifier order.
    #[error("compaction selections are not deterministically ordered")]
    SelectionOrder,
    /// Fail-closed pinned context was omitted.
    #[error("required compaction selection was omitted: {0}")]
    RequiredSelectionOmitted(String),
    /// A memory selection lacks a lifecycle state.
    #[error("memory selection lacks lifecycle state: {0}")]
    MemoryLifecycleMissing(String),
    /// A non-memory selection carries an unrelated lifecycle state.
    #[error("non-memory selection carries a memory lifecycle state: {0}")]
    UnexpectedMemoryLifecycle(String),
    /// Receipt plan digest does not match the embedded plan.
    #[error("compaction receipt plan digest is invalid")]
    PlanDigest,
    /// Receipt input or workspace binding disagrees with its plan.
    #[error("compaction receipt does not match its plan bindings")]
    PlanBinding,
    /// Receipt retained/omitted identifier sets disagree with the plan.
    #[error("compaction receipt selection sets are inconsistent")]
    SelectionSets,
}

#[cfg(test)]
mod tests {
    use serde_json::json;

    use super::*;
    use crate::hash_bytes;

    fn valid_plan() -> CompactionPlan {
        CompactionPlan {
            schema_version: SCHEMA_VERSION.to_owned(),
            algorithm_version: COMPACTION_ALGORITHM_VERSION.to_owned(),
            input_hash: hash_bytes(b"input"),
            contract_hash: hash_bytes(b"contract"),
            workspace_generation: 4,
            state_binding: hash_bytes(b"workspace"),
            side_effect_watermark: 0,
            authority_hash: Some(hash_bytes(b"authority")),
            budget: CompactionBudget::new(100, 10, 1_000).unwrap(),
            selections: vec![ContextSelection {
                id: "contract:task".to_owned(),
                kind: CompactionContextKind::Contract,
                content_hash: hash_bytes(b"selection"),
                decision: CompactionDecision::Retained,
                reason: CompactionSelectionReason::TaskContract,
                priority: CompactionPriority::Pinned,
                required: true,
                rank: 0,
                estimated_tokens: 4,
                byte_count: 9,
                source: None,
                memory_status: None,
            }],
        }
    }

    fn valid_receipt() -> CompactionReceipt {
        let plan = valid_plan();
        CompactionReceipt {
            schema_version: SCHEMA_VERSION.to_owned(),
            algorithm_version: COMPACTION_ALGORITHM_VERSION.to_owned(),
            plan_hash: plan.digest().unwrap(),
            input_hash: plan.input_hash.clone(),
            output_hash: hash_bytes(b"provider-input"),
            output_history_hash: hash_bytes(b"history"),
            output_history_items: 1,
            workspace_generation: plan.workspace_generation,
            state_binding: plan.state_binding.clone(),
            plan,
            usage: CompactionUsage {
                token_estimator: COMPACTION_TOKEN_ESTIMATOR.to_owned(),
                instruction_tokens: 1,
                tool_tokens: 1,
                contract_tokens: 1,
                evidence_tokens: 1,
                memory_tokens: 1,
                history_tokens: 1,
                input_tokens: 6,
                input_bytes: 18,
            },
            retained_ids: vec!["contract:task".to_owned()],
            omitted_ids: Vec::new(),
        }
    }

    #[test]
    fn required_selection_cannot_be_omitted() {
        let mut plan = valid_plan();
        plan.selections[0].decision = CompactionDecision::Omitted;
        assert_eq!(
            plan.validate(),
            Err(CompactionError::RequiredSelectionOmitted(
                "contract:task".to_owned()
            ))
        );
    }

    #[test]
    fn receipt_rejects_plan_tampering_and_commits_history_digest() {
        let mut plan_tampered = valid_receipt();
        plan_tampered.plan.side_effect_watermark = 1;
        assert_eq!(plan_tampered.validate(), Err(CompactionError::PlanDigest));

        let mut history_tampered = valid_receipt();
        history_tampered.output_history_hash = "0".repeat(64);
        assert!(history_tampered.validate().is_ok());
        assert_ne!(
            history_tampered.output_history_hash,
            hash_bytes(b"history"),
            "checkpoint recovery must compare this digest to the persisted history"
        );
    }

    #[test]
    fn default_usage_names_the_versioned_estimator() {
        let usage = CompactionUsage::default();
        assert_eq!(usage.token_estimator, COMPACTION_TOKEN_ESTIMATOR);
        usage.validate(&valid_plan().budget).unwrap();
    }

    #[test]
    fn public_receipt_shape_is_camel_case_and_versioned() {
        let receipt = valid_receipt();
        receipt.validate().unwrap();
        let value = serde_json::to_value(receipt).unwrap();
        assert_eq!(value["schemaVersion"], json!(SCHEMA_VERSION));
        assert_eq!(
            value["algorithmVersion"],
            json!(COMPACTION_ALGORITHM_VERSION)
        );
        assert_eq!(
            value["usage"]["tokenEstimator"],
            json!(COMPACTION_TOKEN_ESTIMATOR)
        );
        assert!(value.get("outputHistoryHash").is_some());
        assert!(value.get("retainedIds").is_some());
        assert!(value.get("output_history_hash").is_none());
    }
}