antimatter 2.0.13

use crate::capsule::common::{CapsuleTag, PolicyDecision, SpanTag};
use crate::session::SessionError;

use crate::opawasm::{DefaultContext, Policy, Runtime};
use async_trait::async_trait;
use itertools::Itertools;
use lru::LruCache;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
use std::num::NonZeroUsize;
use wasmtime::{AsContextMut, Config, Engine, Module, Store};

const ANTIMATTER_ENTRYPOINT: &str = "antimatter/bundle/result";
const CACHE_SIZE: usize = 2000;

/// PolicyRequest represents a policy request that is compatible with the policy engine.
/// The format is:
///  {
///     "tags": {
///       "name": "value",
///       "name": "value",
///       ...
///     },
///     "readParameters": {
///       "name": "value",
///       "name": "value",
///       ...
///     },
///     "domainIdentity": {
///       "name": "value",
///       "name": "value",
///       ...
///     },
///   }
pub type PolicyRequest = HashMap<String, HashMap<String, String>>;

#[derive(Clone, Serialize, Deserialize, PartialEq, Debug)]
/// Represents a section of a data element to which an associated policy applies.
///
/// A `Span` defines a specific range within a data element where a policy is applicable. It includes
/// the starting and ending positions within the data, along with a `policy_request` that specifies
/// the policy-related information for this span.
///
/// # Fields
///
/// * `start` - The starting position of the span within the data element.
/// * `end` - The ending position of the span within the data element.
/// * `policy_request` - A `PolicyRequest` specifying policy-related information for this span.
pub struct Span {
    pub start: usize,
    pub end: usize,
    pub whole_tags: Vec<CapsuleTag>,
    pub policy_request: PolicyRequest,
}

pub struct PolicyEngine {
    policy: Policy<DefaultContext>,
    store: Store<()>,
    // a basic cache for recording policy outcomes
    cache: LruCache<Vec<u8>, PolicyDecision>,
}

#[derive(Clone, PartialEq)]
struct TagEvent {
    index: usize,
    tags_to_add: Vec<CapsuleTag>,
    tags_to_remove: Vec<CapsuleTag>,
}

#[derive(Clone, Serialize, Deserialize, PartialEq, Debug)]
struct RawPolicyResults {
    result: Vec<Vec<Vec<String>>>,
}

#[async_trait]
pub trait PolicyEngineEvaluator {
    /// Evaluates policy decisions for a list of policy input requests using the configured
    /// policy engine.
    ///
    /// # Arguments
    ///
    /// * `policy_inputs` - A vector of policy request references to evaluate.
    ///
    /// # Returns
    ///
    /// A `Result` containing a vector of `PolicyDecision` representing the outcomes of policy
    /// evaluations or an error if the evaluation process fails.
    async fn evaluate(
        &mut self,
        spans: &[Span],
        redact_tags: &[CapsuleTag],
    ) -> Result<Vec<PolicyDecision>, SessionError>;
}

impl PartialOrd for PolicyDecision {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        self.as_int().partial_cmp(&other.as_int())
    }
}

impl PolicyDecision {
    fn as_int(&self) -> u16 {
        // The ordering should match the desired ordering of rules, with
        // a decision's value being greater if the decision is "more
        // restrictive" than the decisions with lower values.
        match self {
            PolicyDecision::NoMatch => 0,
            PolicyDecision::Allow => 1,
            PolicyDecision::Tokenize => 2,
            PolicyDecision::Redact => 3,
            PolicyDecision::DenyRecord => 4,
            PolicyDecision::DenyCapsule => 5,
        }
    }

    fn from_str(s: &str) -> Result<PolicyDecision, SessionError> {
        match s {
            "Allow" => Ok(PolicyDecision::Allow),
            "Redact" => Ok(PolicyDecision::Redact),
            "Tokenize" => Ok(PolicyDecision::Tokenize),
            "DenyRecord" => Ok(PolicyDecision::DenyRecord),
            "DenyCapsule" => Ok(PolicyDecision::DenyCapsule),
            "NoMatch" => Ok(PolicyDecision::NoMatch),
            _ => Err(SessionError::OPAError(format!("unknown policy type {}", s))),
        }
    }

    /// Generates a `PolicyDecision` from a list of policy decisions made by the policy engine.
    ///
    /// # Arguments
    ///
    /// * `input_strings` - A vector of strings representing policy decisions.
    ///
    /// # Returns
    ///
    /// A `Result` containing the resulting `PolicyDecision` based on the list of decisions or an
    /// error if no actionable decision is found.
    ///
    /// This method processes a list of policy decisions in the order they are presented and returns
    /// the first actionable decision encountered. If no actionable decision is found, it returns
    /// `PolicyDecision::NoMatch`.
    pub fn from_str_vec(input_strings: &Vec<String>) -> Result<PolicyDecision, SessionError> {
        for input in input_strings {
            let decision = PolicyDecision::from_str(input)?;
            if decision != PolicyDecision::NoMatch {
                return Ok(decision);
            }
        }
        Ok(PolicyDecision::NoMatch)
    }
}

impl Span {
    fn new(start: usize) -> Self {
        Span {
            start,
            end: 0, // Placeholder, will be set later
            whole_tags: Vec::new(),
            policy_request: HashMap::from([
                ("tags".to_string(), HashMap::new()),
                ("readParameters".to_string(), HashMap::new()),
                ("domainIdentity".to_string(), HashMap::new()),
            ]),
        }
    }
    // A variant of new that includes an end. This is used in testing
    #[cfg(test)]
    fn new_end(start: usize, end: usize) -> Self {
        let mut span = Span::new(start);
        span.end = end;
        span
    }
}

impl TagEvent {
    fn new(index: usize) -> Self {
        TagEvent {
            index,
            tags_to_add: Vec::new(),
            tags_to_remove: Vec::new(),
        }
    }
}

impl PolicyEngine {
    pub async fn new(policy_bundle: &Vec<u8>) -> Result<Self, SessionError> {
        // Configure the WASM runtime
        let mut config = Config::new();
        config.async_support(true);

        let engine =
            Engine::new(&config).map_err(|e| SessionError::OPAError(format!("engine: {}", e)))?;

        // Load the policy WASM module
        let module = Module::new(&engine, policy_bundle)
            .map_err(|e| SessionError::OPAError(format!("failed to load bundle: {}", e)))?;

        // Create a store which will hold the module instance
        let mut store = Store::new(&engine, ());

        // Instantiate the module
        let runtime = Runtime::new(&mut store, &module)
            .await
            .map_err(|e| SessionError::OPAError(format!("runtime: {}", e)))?;

        let policy = runtime
            .without_data(&mut store)
            .await
            .map_err(|e| SessionError::OPAError(format!("policy: {}", e)))?;

        Ok(Self {
            policy,
            store,
            cache: LruCache::new(NonZeroUsize::new(CACHE_SIZE).unwrap()),
        })
    }

    // evaluate_span is an internal function to be caller by the evaluate
    // implementation. If the policy request associated with the span has been
    // evaluated before, the cached result is returned otherwise it is evaluated
    // and the result is cached before being returned.
    async fn evaluate_span(
        &mut self,
        span: &Span,
        redact_tags: &[CapsuleTag],
    ) -> Result<PolicyDecision, SessionError> {
        let cache_key = serde_json::to_vec(&(&span.policy_request, redact_tags)).unwrap();
        let lookup = self.cache.get(&cache_key);
        if let Some(result) = lookup {
            return Ok(*result);
        }

        for tag in &span.whole_tags {
            if redact_tags.contains(tag) {
                self.cache.put(cache_key, PolicyDecision::Redact);
                return Ok(PolicyDecision::Redact);
            }
        }

        let raw_evaluation_result: Value = self
            .policy
            .evaluate(
                self.store.as_context_mut(),
                ANTIMATTER_ENTRYPOINT,
                &vec![&span.policy_request],
            )
            .await
            .map_err(|e| {
                SessionError::OPAError(format!("failed to evaluate policy with data: {}", e))
            })?;

        // The output from the policy engine is a JSON object in the following format:
        // [{"result": [["outcome", ...], ...]}]
        // The outer list contains only one element that has a single field: result. This result
        // is itself an array with an element for each entry in the data_policies parameter. Each
        // element is, in turn, a list of strings containing the outcome of evaluating each rule
        // associated with the template used to generate the policy engine. Consume this and
        // process it into a list of PolicyDecisions
        let converted: Vec<RawPolicyResults> =
            serde_json::from_value(raw_evaluation_result.clone()).unwrap();

        // While the policy returns an array, we expect it to only
        // have one element for now
        if converted.len() != 1 {
            return Err(SessionError::OPAError(format!(
                "malformed evaluation output, expected 1 response, got {}",
                converted.len()
            )));
        }

        let raw_result = converted.first().unwrap();
        let mut output: PolicyDecision = PolicyDecision::NoMatch;
        for input in raw_result.result.iter() {
            for policy in input.iter() {
                for rule in policy.iter() {
                    let parts = rule.split(';').collect::<Vec<_>>();
                    let decision_str = match parts.get(2) {
                        Some(s) => s,
                        None => {
                            return Err(SessionError::OPAError(format!(
                                "malformed policy decision {}",
                                rule,
                            )))
                        }
                    };

                    let decision = PolicyDecision::from_str(decision_str)?;

                    if decision != PolicyDecision::NoMatch {
                        // across all policies, choose the most restrictive
                        // (i.e. "greatest") decision.
                        if decision > output {
                            output = decision;
                        }
                        // within a policy, only take the first decision
                        // that is not NoMatch.
                        break;
                    }
                }
            }
        }
        Ok(output)
    }
}

#[async_trait]
impl PolicyEngineEvaluator for PolicyEngine {
    async fn evaluate(
        &mut self,
        spans: &[Span],
        redact_tags: &[CapsuleTag],
    ) -> Result<Vec<PolicyDecision>, SessionError> {
        let mut result: Vec<PolicyDecision> = Vec::new();
        for span in spans.iter() {
            result.push(self.evaluate_span(span, redact_tags).await?)
        }
        Ok(result)
    }
}

/// Generates non-overlapping spans to fully cover a data element based on provided span tags,
/// read parameters, domain identity, and data length.
///
/// # Arguments
///
/// * `span_tags` - A list of span tags associated with the data element.
/// * `column_tags` - A list of column tags associated with the data element.
/// * `read_parameters` - A set of key-value pairs to be included in policy requests.
/// * `capsule_tags` - A list of capsule tags associated with the data element.
/// * `domain_identity` - A set of key-value pairs to be included in policy requests.
/// * `data_length` - The length of the data element.
///
/// # Returns
///
/// A vector of `Span` objects representing the computed spans that cover the entire data
/// element. Each span contains a policy request as well as references to the the start and end
/// of the data segment it is associated with.
pub fn generate_spans(
    span_tags: &[SpanTag],
    column_tags: &[CapsuleTag],
    read_parameters: &HashMap<String, String>,
    capsule_tags: &[CapsuleTag],
    domain_identity: &HashMap<String, String>,
    data_length: usize,
) -> Vec<Span> {
    // Begin by processing the tags into an ordered list of tag events. These
    // will dictate the boundaries between policy requests ad read parameters
    // and domain identities span the whole element.

    let mut tag_events: HashMap<usize, TagEvent> = HashMap::new();

    for span_tag in span_tags {
        tag_events
            .entry(span_tag.start)
            .or_insert_with(|| TagEvent::new(span_tag.start))
            .tags_to_add
            .push(span_tag.tag.clone());
        tag_events
            .entry(span_tag.end)
            .or_insert_with(|| TagEvent::new(span_tag.end))
            .tags_to_remove
            .push(span_tag.tag.clone());
    }

    // Next, use the tag events to generate an ordered list of spans that
    // covert the data element
    let mut spans: Vec<Span> = Vec::new();
    // populate tags with the capsule_tags and column tags as these should
    // appear in every span
    let mut init_tags: Vec<CapsuleTag> = capsule_tags.to_owned();
    init_tags.extend(column_tags.to_owned());
    let mut current_span: Span = Span::new(0);
    // tail_boundary indicates the last boundary index observed. The difference
    // between this and data_length indicates how much of the data element the
    // tail span should cover.
    let mut tail_boundary = 0;

    // we need an initial tag to support fill resource tags on the first element
    tag_events
        .entry(0)
        .or_insert_with(|| TagEvent::new(0))
        .tags_to_add
        .extend(init_tags.clone());

    let mut tags: Vec<CapsuleTag> = Vec::new();
    let mut tag_count: HashMap<CapsuleTag, usize> = HashMap::new(); // New: Track tag counts
    for key in tag_events.keys().sorted() {
        // A new tag event indicates a boundary, close out the current span
        // and record it (provided it spans at least one byte). Next apply
        // the boundary change and start a new span.
        current_span.end = *key;
        if current_span.start < current_span.end {
            current_span
                .policy_request
                .insert("readParameters".to_string(), read_parameters.clone());
            current_span
                .policy_request
                .insert("domainIdentity".to_string(), domain_identity.clone());
            spans.push(current_span.clone());
        }
        let event = &tag_events[key];
        // Modified: Update tag counts and manage tags
        for tag in &event.tags_to_add {
            *tag_count.entry(tag.clone()).or_insert(0) += 1;
            if tag_count[tag] == 1 {
                tags.push(tag.clone());
            }
        }
        for tag in &event.tags_to_remove {
            if let Some(count) = tag_count.get_mut(tag) {
                *count -= 1;
                if *count == 0 {
                    tags.retain(|t| t != tag);
                }
            }
        }
        current_span = Span::new(*key);
        for tag in tags.iter_mut() {
            let current_span_tags = current_span.policy_request.get_mut("tags").unwrap();
            current_span_tags.insert(tag.name.clone().to_string(), tag.value.clone().to_string());
            current_span.whole_tags.push(tag.clone());
        }

        tail_boundary = *key;
    }

    // Finally create a span for the remaining data (if needed). As all tag
    // event have resolved, this span has no tags aside from the capsule tags.
    if tail_boundary < data_length {
        let mut tail = Span::new(tail_boundary);
        // Note: use data_length and not (data_length - 1) as the end is exclusive.
        tail.end = data_length;
        tail.policy_request
            .insert("readParameters".to_string(), read_parameters.clone());
        tail.policy_request
            .insert("domainIdentity".to_string(), domain_identity.clone());
        for tag in tags.iter_mut() {
            let tail_tags = tail.policy_request.get_mut("tags").unwrap();
            tail_tags.insert(tag.name.clone().to_string(), tag.value.clone().to_string());
            tail.whole_tags.push(tag.clone());
        }
        spans.push(tail);
    }

    spans
}

/// Enforces policies on a given data element based on provided spans.
///
/// This function accepts a list of spans and the associated data element. It then evaluates
/// and applies the policy operation specified for each span. If a "DenyRecord" conclusion is
/// reached, the original data is returned unmodified along with a "DenyRecord" PolicyDecision.
/// Otherwise, the modified data is returned along with a "NoMatch" PolicyDecision.
///
/// # Arguments
///
/// * `data` - A reference to the data element on which policies will be enforced.
/// * `spans` - A vector of `Span` objects representing sections of the data element with
///            associated policies.
/// * `engine` - A mutable reference to a policy engine that implements the
///             `PolicyEngineEvaluator` trait.
///
/// # Returns
///
/// A `Result` containing a tuple of the modified data as a `String`, `Vec<SpanTag>` and the
/// resulting `PolicyDecision`.
///
/// # Errors
///
/// Returns a `SessionError` if there is a mismatch between the number of spans and policy
/// decisions or if an error occurs during policy evaluation.
///
/// # TODO
///
/// Introduce support for tokenization.
pub async fn enforce_policies<T>(
    data: &[u8],
    spans: Vec<Span>,
    redact_tags: &[CapsuleTag],
    engine: &mut T,
) -> Result<(Vec<u8>, Vec<(PolicyDecision, SpanTag)>, PolicyDecision), SessionError>
where
    T: PolicyEngineEvaluator,
{
    let mut adjusted_span_tags: Vec<(PolicyDecision, SpanTag)> = Vec::new();

    let policy_decisions = engine.evaluate(&spans, redact_tags).await?;

    // Note: the spans must match the policy_decisions, otherwise something has gone horribly wrong
    if spans.len() != policy_decisions.len() {
        return Err(SessionError::OPAError(format!(
            "length mismatch between spans ({}) and decisions ({})",
            spans.len(),
            policy_decisions.len()
        )));
    }

    let mut result: Vec<u8> = vec![];
    let mut previously_redacted = false;

    for (span, decision) in spans.iter().zip(policy_decisions.iter()) {
        if *decision == PolicyDecision::DenyRecord {
            // TODO: why copy data here when it should (?) all be redacted
            return Ok((
                data.to_owned(),
                adjusted_span_tags,
                PolicyDecision::DenyRecord,
            ));
        }
        if *decision == PolicyDecision::DenyCapsule {
            // TODO: why copy data here when it should (?) all be redacted
            return Ok((
                data.to_owned(),
                adjusted_span_tags,
                PolicyDecision::DenyCapsule,
            ));
        }
        let span_result: Vec<u8> = match decision {
            PolicyDecision::NoMatch | PolicyDecision::Tokenize | PolicyDecision::Allow => {
                //TODO: Implement Tokenize.
                previously_redacted = false;
                if span.start > span.end || span.end > data.len() {
                    return Err(SessionError::OPAError(format!(
                        "Span is out of bounds. Data element: [{}:{}], span: [{}:{}]",
                        0,
                        data.len(),
                        span.start,
                        span.end
                    )));
                }
                Ok(data[span.start..span.end].to_vec())
            }
            PolicyDecision::Redact => {
                if previously_redacted {
                    Ok("".as_bytes().to_vec())
                } else {
                    previously_redacted = true;
                    Ok("{redacted}".as_bytes().to_vec())
                }
            }
            PolicyDecision::DenyRecord => Err(SessionError::OPAError(
                "attempted to match a DenyRecord decision".to_string(),
            )), // should not be possible
            PolicyDecision::DenyCapsule => Err(SessionError::OPAError(
                "attempted to match a DenyCapsule decision".to_string(),
            )), // should not be possible
        }?;

        let before = result.len();
        result = [result.as_slice(), span_result.as_slice()].concat();

        // Generate span tags for rendered output.
        for whole_tag in span.whole_tags.iter() {
            adjusted_span_tags.push((
                *decision,
                SpanTag {
                    start: before,
                    end: result.len(),
                    tag: whole_tag.clone(),
                },
            ));
        }
    }

    Ok((result, adjusted_span_tags, PolicyDecision::NoMatch))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::capsule::common::TagType;

    #[test]
    fn test_policy_decision_from_str() {
        // Test valid decision strings
        assert_eq!(
            PolicyDecision::from_str("Allow").unwrap(),
            PolicyDecision::Allow
        );
        assert_eq!(
            PolicyDecision::from_str("Redact").unwrap(),
            PolicyDecision::Redact
        );
        assert_eq!(
            PolicyDecision::from_str("Tokenize").unwrap(),
            PolicyDecision::Tokenize
        );
        assert_eq!(
            PolicyDecision::from_str("DenyRecord").unwrap(),
            PolicyDecision::DenyRecord
        );
        assert_eq!(
            PolicyDecision::from_str("DenyCapsule").unwrap(),
            PolicyDecision::DenyCapsule
        );
        assert_eq!(
            PolicyDecision::from_str("NoMatch").unwrap(),
            PolicyDecision::NoMatch
        );

        // Test invalid decision string
        assert!(matches!(
            PolicyDecision::from_str("InvalidDecision"),
            Err(SessionError::OPAError(_))
        ));
    }

    #[test]
    fn test_policy_decision_from_str_vec() {
        // Test a list with all "NoMatch" decisions
        let no_match_decisions = vec![
            "NoMatch".to_string(),
            "NoMatch".to_string(),
            "NoMatch".to_string(),
        ];
        assert_eq!(
            PolicyDecision::from_str_vec(&no_match_decisions).unwrap(),
            PolicyDecision::NoMatch
        );

        // Test an empty decision list
        let empty_decisions: Vec<String> = Vec::new();
        assert_eq!(
            PolicyDecision::from_str_vec(&empty_decisions).unwrap(),
            PolicyDecision::NoMatch
        );

        // Test a list with a mix of "NoMatch" and other decisions
        let mixed_decisions = vec![
            "NoMatch".to_string(),
            "NoMatch".to_string(),
            "Redact".to_string(),
            "NoMatch".to_string(),
            "Tokenize".to_string(),
        ];
        assert_eq!(
            PolicyDecision::from_str_vec(&mixed_decisions).unwrap(),
            PolicyDecision::Redact
        );
    }

    fn tag_a() -> CapsuleTag {
        CapsuleTag {
            name: "test1".to_string(),
            tag_type: TagType::Unary,
            value: "none".to_string(),
            source: "src".to_string(),
            hook_version: (0, 0, 0),
        }
    }

    fn tag_b() -> CapsuleTag {
        CapsuleTag {
            name: "test2".to_string(),
            tag_type: TagType::Unary,
            value: "none".to_string(),
            source: "src".to_string(),
            hook_version: (0, 0, 0),
        }
    }

    #[test]
    fn test_non_overlapping_tags() {
        let span_tags = vec![
            SpanTag {
                start: 0,
                end: 11,
                tag: tag_a(),
            },
            SpanTag {
                start: 11,
                end: 21,
                tag: tag_b(),
            },
        ];
        let read_parameters = HashMap::new();
        let domain_identity = HashMap::new();
        let data_length = 21;
        let spans = generate_spans(
            &span_tags,
            &vec![],
            &read_parameters,
            &vec![],
            &domain_identity,
            data_length,
        );
        assert_eq!(spans.len(), 2);
        assert_eq!(
            spans.get(0).unwrap().policy_request,
            HashMap::from([
                (
                    "tags".to_string(),
                    HashMap::from([("test1".to_string(), "none".to_string())])
                ),
                ("readParameters".to_string(), HashMap::new()),
                ("domainIdentity".to_string(), HashMap::new())
            ])
        );
        assert_eq!(
            spans.get(1).unwrap().policy_request,
            HashMap::from([
                (
                    "tags".to_string(),
                    HashMap::from([("test2".to_string(), "none".to_string())])
                ),
                ("readParameters".to_string(), HashMap::new()),
                ("domainIdentity".to_string(), HashMap::new())
            ])
        );
    }

    #[test]
    fn test_overlapping_tags() {
        let span_tags = vec![
            SpanTag {
                start: 0,
                end: 15,
                tag: tag_a(),
            },
            SpanTag {
                start: 10,
                end: 21,
                tag: tag_b(),
            },
        ];
        let read_parameters = HashMap::new();
        let domain_identity = HashMap::new();
        let data_length = 21;

        let spans = generate_spans(
            &span_tags,
            &vec![],
            &read_parameters,
            &vec![],
            &domain_identity,
            data_length,
        );
        assert_eq!(spans.len(), 3);
        assert_eq!(
            spans.get(0).unwrap().policy_request,
            HashMap::from([
                (
                    "tags".to_string(),
                    HashMap::from([("test1".to_string(), "none".to_string())])
                ),
                ("readParameters".to_string(), HashMap::new()),
                ("domainIdentity".to_string(), HashMap::new())
            ])
        );
        assert_eq!(
            spans.get(1).unwrap().policy_request,
            HashMap::from([
                (
                    "tags".to_string(),
                    HashMap::from([
                        ("test1".to_string(), "none".to_string()),
                        ("test2".to_string(), "none".to_string())
                    ])
                ),
                ("readParameters".to_string(), HashMap::new()),
                ("domainIdentity".to_string(), HashMap::new())
            ])
        );
        assert_eq!(
            spans.get(2).unwrap().policy_request,
            HashMap::from([
                (
                    "tags".to_string(),
                    HashMap::from([("test2".to_string(), "none".to_string())])
                ),
                ("readParameters".to_string(), HashMap::new()),
                ("domainIdentity".to_string(), HashMap::new())
            ])
        );
        assert_eq!(spans.get(2).unwrap().end, 21);
    }

    struct MockPolicyEngine {
        responses: Vec<PolicyDecision>,
    }

    #[async_trait::async_trait]
    impl PolicyEngineEvaluator for MockPolicyEngine {
        async fn evaluate(
            &mut self,
            pr: &[Span],
            _redact_tags: &[CapsuleTag],
        ) -> Result<Vec<PolicyDecision>, SessionError> {
            Ok(pr
                .iter()
                .enumerate()
                .map(|(idx, &_)| self.responses.get(idx).unwrap().clone())
                .collect())
        }
    }

    #[tokio::test]
    async fn test_enforce_policies_allow() {
        let data = "This is some data".as_bytes().to_vec();
        let spans = vec![
            Span::new_end(0, 5),
            Span::new_end(5, 10),
            Span::new_end(10, data.len()),
        ];

        let mut engine = MockPolicyEngine {
            responses: vec![
                PolicyDecision::Allow,
                PolicyDecision::Allow,
                PolicyDecision::Allow,
            ],
        };
        let (result, span_tags, decision) =
            enforce_policies(&data, spans, &Vec::new(), &mut engine)
                .await
                .unwrap();

        assert_eq!(result, data);
        assert_eq!(span_tags, vec![]);
        assert_eq!(decision, PolicyDecision::NoMatch);
    }

    #[tokio::test]
    async fn test_enforce_policies_redact() {
        let data = "This is some data".as_bytes().to_vec();
        let spans = vec![
            Span::new_end(0, 5),
            Span::new_end(5, 7),
            Span::new_end(7, data.len()),
        ];

        let mut engine = MockPolicyEngine {
            responses: vec![
                PolicyDecision::Allow,
                PolicyDecision::Redact,
                PolicyDecision::Allow,
            ],
        };
        let (result, span_tags, decision) =
            enforce_policies(&data, spans, &Vec::new(), &mut engine)
                .await
                .unwrap();

        assert_eq!(result, "This {redacted} some data".as_bytes().to_vec());
        assert_eq!(span_tags, vec![]);
        assert_eq!(decision, PolicyDecision::NoMatch);
    }

    #[tokio::test]
    async fn test_enforce_policies_deny_row() {
        let data = "This is some data".as_bytes().to_vec();
        let spans = vec![
            Span::new_end(0, 5),
            Span::new_end(5, 7),
            Span::new_end(7, data.len()),
        ];

        let mut engine = MockPolicyEngine {
            responses: vec![
                PolicyDecision::Allow,
                PolicyDecision::Redact,
                PolicyDecision::DenyRecord,
            ],
        };
        let (result, span_tags, decision) =
            enforce_policies(&data, spans, &Vec::new(), &mut engine)
                .await
                .unwrap();

        assert_eq!(result, data);
        assert_eq!(span_tags, vec![]);
        assert_eq!(decision, PolicyDecision::DenyRecord);
    }

    #[tokio::test]
    async fn test_enforce_policies_deny_capsule() {
        let data = "This is some data".as_bytes().to_vec();
        let spans = vec![
            Span::new_end(0, 5),
            Span::new_end(5, 7),
            Span::new_end(7, data.len()),
        ];

        let mut engine = MockPolicyEngine {
            responses: vec![
                PolicyDecision::Allow,
                PolicyDecision::Redact,
                PolicyDecision::DenyCapsule,
            ],
        };
        let (result, span_tags, decision) =
            enforce_policies(&data, spans, &Vec::new(), &mut engine)
                .await
                .unwrap();

        assert_eq!(result, data);
        assert_eq!(span_tags, vec![]);
        assert_eq!(decision, PolicyDecision::DenyCapsule);
    }

    #[tokio::test]
    async fn test_enforce_policies_out_of_bounds() {
        let data = "This is some data".as_bytes().to_vec();
        let spans = vec![Span::new_end(2, data.len() + 1)];

        let mut engine = MockPolicyEngine {
            responses: vec![PolicyDecision::Allow],
        };
        let outcome = enforce_policies(&data, spans, &Vec::new(), &mut engine).await;

        assert_eq!(outcome.is_err(), true);
        assert_eq!(
            outcome.err().unwrap(),
            SessionError::OPAError(
                "Span is out of bounds. Data element: [0:17], span: [2:18]".to_string()
            )
        );
    }
    #[test]
    fn test_overlapping_same_tags() {
        let span_tags = vec![
            SpanTag {
                start: 1434,
                end: 1436,
                tag: CapsuleTag {
                    name: "tag.antimatter.io/pii/financial_terms".to_string(),
                    tag_type: TagType::Unary,
                    value: "".to_string(),
                    source: "custom".to_string(),
                    hook_version: (1, 0, 0),
                },
            },
            SpanTag {
                start: 1434,
                end: 1442,
                tag: CapsuleTag {
                    name: "tag.antimatter.io/pii/financial_terms".to_string(),
                    tag_type: TagType::Unary,
                    value: "".to_string(),
                    source: "custom".to_string(),
                    hook_version: (1, 0, 0),
                },
            },
        ];
        let read_parameters = HashMap::new();
        let domain_identity = HashMap::new();
        let data_length = 1500;
        let capsule_tags = vec![CapsuleTag {
            name: "tag.antimatter.io/pii/cap/financial_terms".to_string(),
            tag_type: TagType::Unary,
            value: "".to_string(),
            source: "custom".to_string(),
            hook_version: (1, 0, 0),
        }];

        let spans = generate_spans(
            &span_tags,
            &vec![],
            &read_parameters,
            &capsule_tags,
            &domain_identity,
            data_length,
        );

        assert_eq!(spans.len(), 4);
        assert_eq!(spans.get(0).unwrap().whole_tags.len(), 1);
        assert_eq!(
            spans.get(0).unwrap().whole_tags.get(0),
            Some(&CapsuleTag {
                name: "tag.antimatter.io/pii/cap/financial_terms".to_string(),
                tag_type: TagType::Unary,
                value: "".to_string(),
                source: "custom".to_string(),
                hook_version: (1, 0, 0)
            })
        );
        assert_eq!(spans.get(1).unwrap().whole_tags.len(), 2);
        assert_eq!(
            spans.get(1).unwrap().whole_tags.get(0),
            Some(&CapsuleTag {
                name: "tag.antimatter.io/pii/cap/financial_terms".to_string(),
                tag_type: TagType::Unary,
                value: "".to_string(),
                source: "custom".to_string(),
                hook_version: (1, 0, 0)
            })
        );
        assert_eq!(
            spans.get(1).unwrap().whole_tags.get(1),
            Some(&CapsuleTag {
                name: "tag.antimatter.io/pii/financial_terms".to_string(),
                tag_type: TagType::Unary,
                value: "".to_string(),
                source: "custom".to_string(),
                hook_version: (1, 0, 0)
            })
        );
        assert_eq!(spans.get(2).unwrap().whole_tags.len(), 2);
        assert_eq!(
            spans.get(2).unwrap().whole_tags.get(0),
            Some(&CapsuleTag {
                name: "tag.antimatter.io/pii/cap/financial_terms".to_string(),
                tag_type: TagType::Unary,
                value: "".to_string(),
                source: "custom".to_string(),
                hook_version: (1, 0, 0)
            })
        );
        assert_eq!(
            spans.get(2).unwrap().whole_tags.get(1),
            Some(&CapsuleTag {
                name: "tag.antimatter.io/pii/financial_terms".to_string(),
                tag_type: TagType::Unary,
                value: "".to_string(),
                source: "custom".to_string(),
                hook_version: (1, 0, 0)
            })
        );
        assert_eq!(spans.get(3).unwrap().whole_tags.len(), 1);
        assert_eq!(
            spans.get(3).unwrap().whole_tags.get(0),
            Some(&CapsuleTag {
                name: "tag.antimatter.io/pii/cap/financial_terms".to_string(),
                tag_type: TagType::Unary,
                value: "".to_string(),
                source: "custom".to_string(),
                hook_version: (1, 0, 0)
            })
        );
    }

    fn create_capsule_tag(name: &str) -> CapsuleTag {
        CapsuleTag {
            name: name.to_string(),
            tag_type: TagType::Unary,
            value: "".to_string(),
            source: "custom".to_string(),
            hook_version: (1, 0, 0),
        }
    }

    fn create_span_tag(start: usize, end: usize, name: &str) -> SpanTag {
        SpanTag {
            start,
            end,
            tag: create_capsule_tag(name),
        }
    }

    #[test]
    fn test_multiple_overlapping_spans() {
        let span_tags = vec![
            create_span_tag(100, 200, "tag1"),
            create_span_tag(150, 250, "tag1"),
            create_span_tag(180, 300, "tag1"),
        ];
        let capsule_tags = vec![create_capsule_tag("capsule")];
        let spans = generate_spans(
            &span_tags,
            &[],
            &HashMap::new(),
            &capsule_tags,
            &HashMap::new(),
            400,
        );

        assert_eq!(spans.len(), 7);
        assert_eq!(spans[0].whole_tags.len(), 1); // 0-100: only capsule tag
        assert_eq!(spans[1].whole_tags.len(), 2); // 100-150: capsule + tag1
        assert_eq!(spans[2].whole_tags.len(), 2); // 150-180: capsule + tag1
        assert_eq!(spans[3].whole_tags.len(), 2); // 180-200: capsule + tag1
        assert_eq!(spans[4].whole_tags.len(), 2); // 200-250: capsule + tag1
        assert_eq!(spans[5].whole_tags.len(), 2); // 250-300: capsule + tag1
        assert_eq!(spans[6].whole_tags.len(), 1); // 300-400: only capsule tag
    }

    #[test]
    fn test_spans_starting_and_ending_together() {
        let span_tags = vec![
            create_span_tag(100, 200, "tag1"),
            create_span_tag(100, 200, "tag2"),
            create_span_tag(200, 300, "tag3"),
        ];
        let capsule_tags = vec![create_capsule_tag("capsule")];
        let spans = generate_spans(
            &span_tags,
            &[],
            &HashMap::new(),
            &capsule_tags,
            &HashMap::new(),
            400,
        );

        assert_eq!(spans.len(), 4);
        assert_eq!(spans[0].whole_tags.len(), 1); // 0-100: only capsule tag
        assert_eq!(spans[1].whole_tags.len(), 3); // 100-200: capsule + tag1 + tag2
        assert_eq!(spans[2].whole_tags.len(), 2); // 200-300: capsule + tag3
        assert_eq!(spans[3].whole_tags.len(), 1); // 300-400: only capsule tag
    }

    #[test]
    fn test_span_covering_entire_range() {
        let span_tags = vec![create_span_tag(0, 400, "tag1")];
        let capsule_tags = vec![create_capsule_tag("capsule")];
        let spans = generate_spans(
            &span_tags,
            &[],
            &HashMap::new(),
            &capsule_tags,
            &HashMap::new(),
            400,
        );

        assert_eq!(spans.len(), 1);
        assert_eq!(spans[0].whole_tags.len(), 2); // 0-400: capsule + tag1
    }

    #[test]
    fn test_no_spans() {
        let capsule_tags = vec![create_capsule_tag("capsule")];
        let spans = generate_spans(
            &[],
            &[],
            &HashMap::new(),
            &capsule_tags,
            &HashMap::new(),
            400,
        );

        assert_eq!(spans.len(), 1);
        assert_eq!(spans[0].whole_tags.len(), 1); // 0-400: only capsule tag
    }

    #[test]
    fn test_overlapping_different_tags() {
        let span_tags = vec![
            create_span_tag(100, 300, "tag1"),
            create_span_tag(200, 400, "tag2"),
        ];
        let capsule_tags = vec![create_capsule_tag("capsule")];
        let spans = generate_spans(
            &span_tags,
            &[],
            &HashMap::new(),
            &capsule_tags,
            &HashMap::new(),
            500,
        );

        assert_eq!(spans.len(), 5);
        assert_eq!(spans[0].whole_tags.len(), 1); // 0-100: only capsule tag
        assert_eq!(spans[1].whole_tags.len(), 2); // 100-200: capsule + tag1
        assert_eq!(spans[2].whole_tags.len(), 3); // 200-300: capsule + tag1 + tag2
        assert_eq!(spans[3].whole_tags.len(), 2); // 300-400: capsule + tag2
        assert_eq!(spans[4].whole_tags.len(), 1); // 400-500: only capsule tag
    }

    #[test]
    fn test_zero_length_spans() {
        let span_tags = vec![
            create_span_tag(100, 100, "tag1"),
            create_span_tag(200, 200, "tag2"),
        ];
        let capsule_tags = vec![create_capsule_tag("capsule")];
        let spans = generate_spans(
            &span_tags,
            &[],
            &HashMap::new(),
            &capsule_tags,
            &HashMap::new(),
            300,
        );

        assert_eq!(spans.len(), 3);
        assert_eq!(spans[0].whole_tags.len(), 1); // 0-100: only capsule tag
        assert_eq!(spans[1].whole_tags.len(), 1); // 100-200: only capsule tag
        assert_eq!(spans[2].whole_tags.len(), 1); // 200-300: only capsule tag
    }
}