User Query → PolicyEngine Evaluation → Query Execution → Masking Transform → Response
                  |                                            |
            Based on ABAC                                Based on field
            attributes                                   masking policies
            determine masking level                      apply masking mode

┌─────────────────────────────────────────────────┐
│                   API Gateway                    │
│              (gRPC / REST Entry)                 │
└──────────────────────┬──────────────────────────┘
                       │
┌──────────────────────▼──────────────────────────┐
│              Query Execution Engine              │
│           (Iceberg + Nessie Queries)             │
└──────────────────────┬──────────────────────────┘
                       │ Raw Result Set
┌──────────────────────▼──────────────────────────┐
│            Dynamic Masking Engine                │
│  ┌─────────┐  ┌──────────┐  ┌───────────────┐  │
│  │ Policy   │  │ Field    │  │ Masking       │  │
│  │ Resolver │→│ Classifier│→│ Transformer   │  │
│  └─────────┘  └──────────┘  └───────────────┘  │
│       ↑              ↑              │           │
│  ABAC Engine    Schema Registry    6 Modes      │
└──────────────────────┬──────────────────────────┘
                       │ Masked Result Set
┌──────────────────────▼──────────────────────────┐
│              Response Serializer                 │
│           (gRPC Protobuf Serialization)          │
└─────────────────────────────────────────────────┘

class DynamicMaskingEngine:
    """Dynamic Masking Engine - Query result post-processor"""

    def __init__(
        self,
        policy_resolver: PolicyResolver,
        field_classifier: FieldClassifier,
        masking_registry: MaskingTransformerRegistry,
    ):
        self._policy_resolver = policy_resolver
        self._field_classifier = field_classifier
        self._masking_registry = masking_registry

    async def apply_masking(
        self,
        result_set: ResultSet,
        request_context: RequestContext,
    ) -> ResultSet:
        """Apply dynamic masking to query result set"""
        # 1. Resolve masking policies for current user
        policies = await self._policy_resolver.resolve(
            subject=request_context.subject,
            resource=result_set.source_object,
            environment=request_context.environment,
        )

        # 2. Determine masking mode for each field
        field_masks = {}
        for field in result_set.schema.fields:
            classification = self._field_classifier.classify(field)
            mask_mode = policies.get_mask_mode(classification)
            if mask_mode:
                field_masks[field.name] = mask_mode

        # 3. Apply masking transformations
        masked_result = result_set.copy()
        for row in masked_result.rows:
            for field_name, mode in field_masks.items():
                transformer = self._masking_registry.get(mode)
                row[field_name] = transformer.transform(
                    value=row[field_name],
                    field_name=field_name,
                    context=request_context,
                )

        return masked_result

class PolicyResolver:
    """Masking Policy Resolver"""

    async def resolve(
        self,
        subject: Subject,
        resource: OntologyObject,
        environment: Environment,
    ) -> MaskingPolicies:
        """Resolve applicable masking policies"""
        evaluation = await self._abac_engine.evaluate(
            subject_attrs={
                "roles": subject.roles,
                "clearance_level": subject.clearance_level,
                "department": subject.department,
            },
            resource_attrs={
                "classification": resource.classification,
                "sensitivity": resource.sensitivity_label,
                "data_domain": resource.data_domain,
            },
            environment_attrs={
                "time": environment.request_time,
                "network": environment.network_zone,
                "client_type": environment.client_type,
            },
        )

        return MaskingPolicies.from_evaluation(evaluation)

class FullMaskingTransformer(MaskingTransformer):
    """Full Masking - Replace field values with mask characters"""

    mode = MaskingMode.FULL

    def transform(self, value: Any, field_name: str, context: RequestContext) -> str:
        if value is None:
            return None

        if isinstance(value, str):
            return "***" if len(value) <= 10 else "******"
        elif isinstance(value, (int, float)):
            return 0
        elif isinstance(value, datetime):
            return datetime(1970, 1, 1)
        else:
            return "***"

class PartialMaskingTransformer(MaskingTransformer):
    """Partial Masking - Keep prefix/suffix, mask the middle"""

    mode = MaskingMode.PARTIAL

    def __init__(self, prefix_len: int = 3, suffix_len: int = 4, mask_char: str = "*"):
        self._prefix_len = prefix_len
        self._suffix_len = suffix_len
        self._mask_char = mask_char

    def transform(self, value: Any, field_name: str, context: RequestContext) -> str:
        if value is None:
            return None

        s = str(value)
        if len(s) <= self._prefix_len + self._suffix_len:
            return self._mask_char * len(s)

        prefix = s[:self._prefix_len]
        suffix = s[-self._suffix_len:]
        masked_len = len(s) - self._prefix_len - self._suffix_len
        return f"{prefix}{self._mask_char * masked_len}{suffix}"

FIELD_MASKING_PROFILES = {
    "phone": PartialMaskingTransformer(prefix_len=3, suffix_len=4),
    "email": EmailMaskingTransformer(),       # user → u***r
    "id_card": PartialMaskingTransformer(prefix_len=6, suffix_len=4),
    "bank_card": PartialMaskingTransformer(prefix_len=4, suffix_len=4),
    "name": PartialMaskingTransformer(prefix_len=1, suffix_len=0),
}

class HashReplacementTransformer(MaskingTransformer):
    """Hash Replacement - Generate deterministic substitutes using HMAC-SHA256"""

    mode = MaskingMode.HASH

    def __init__(self, secret_key: bytes, output_format: str = "hex16"):
        self._secret_key = secret_key
        self._output_format = output_format

    def transform(self, value: Any, field_name: str, context: RequestContext) -> str:
        if value is None:
            return None

        # HMAC ensures same input produces same output (supports JOIN)
        mac = hmac.new(
            self._secret_key,
            msg=f"{field_name}:{value}".encode("utf-8"),
            digestmod=hashlib.sha256,
        )

        digest = mac.hexdigest()

        if self._output_format == "hex16":
            return digest[:16]
        elif self._output_format == "uuid":
            return str(uuid.UUID(digest[:32]))
        else:
            return digest

class RangeGeneralizationTransformer(MaskingTransformer):
    """Range Generalization - Replace precise values with ranges"""

    mode = MaskingMode.RANGE

    def __init__(self, ranges: list[tuple] | None = None, step: int | None = None):
        self._ranges = ranges
        self._step = step

    def transform(self, value: Any, field_name: str, context: RequestContext) -> str:
        if value is None:
            return None

        if isinstance(value, (int, float)):
            return self._generalize_number(value)
        elif isinstance(value, datetime):
            return self._generalize_date(value)
        elif isinstance(value, date):
            return self._generalize_date(datetime.combine(value, datetime.min.time()))
        else:
            return str(value)

    def _generalize_number(self, value: float) -> str:
        if self._ranges:
            for low, high in self._ranges:
                if low <= value < high:
                    return f"{low}-{high}"
            return "Other"

        if self._step:
            lower = (value // self._step) * self._step
            upper = lower + self._step
            return f"{int(lower)}-{int(upper)}"

        magnitude = 10 ** max(0, len(str(int(abs(value)))) - 1)
        lower = (value // magnitude) * magnitude
        upper = lower + magnitude
        return f"{int(lower)}-{int(upper)}"

    def _generalize_date(self, value: datetime) -> str:
        return value.strftime("%Y-%m")

class FPETransformer(MaskingTransformer):
    """Format-Preserving Encryption - FF1/FF3-1 algorithm implementation"""

    mode = MaskingMode.FPE

    def __init__(self, key: bytes, tweak: bytes):
        self._cipher = FF3Cipher(key=key.hex(), tweak=tweak.hex(), radix=10)
        self._alpha_cipher = FF3Cipher(key=key.hex(), tweak=tweak.hex(), radix=36)

    def transform(self, value: Any, field_name: str, context: RequestContext) -> str:
        if value is None:
            return None

        s = str(value)

        # Separate format characters from effective characters
        format_map = []
        effective_chars = []
        for i, ch in enumerate(s):
            if ch.isdigit():
                effective_chars.append(ch)
                format_map.append(("digit", i))
            elif ch.isalpha():
                effective_chars.append(ch.lower())
                format_map.append(("alpha", i))
            else:
                format_map.append(("literal", i, ch))

        # Apply FPE encryption to effective characters
        if all(fc[0] == "digit" for fc in format_map if fc[0] != "literal"):
            encrypted = self._cipher.encrypt("".join(effective_chars))
        else:
            encrypted = self._alpha_cipher.encrypt("".join(effective_chars))

        # Reconstruct with original format
        result = list(s)
        enc_idx = 0
        for fm in format_map:
            if fm[0] in ("digit", "alpha"):
                result[fm[1]] = encrypted[enc_idx]
                enc_idx += 1

        return "".join(result)

class ConditionalMaskingTransformer(MaskingTransformer):
    """Conditional Masking - Dynamically select strategy based on context"""

    mode = MaskingMode.CONDITIONAL

    def __init__(self, rules: list[ConditionalRule]):
        self._rules = rules

    def transform(self, value: Any, field_name: str, context: RequestContext) -> Any:
        if value is None:
            return None

        for rule in self._rules:
            if rule.evaluate(context):
                return rule.transformer.transform(value, field_name, context)

        # No matching rule: fall back to full masking (safe default)
        return FullMaskingTransformer().transform(value, field_name, context)

phone_masking = ConditionalMaskingTransformer(rules=[
    # Rule 1: Data owner sees full value
    ConditionalRule(
        condition=lambda ctx: ctx.subject.id == ctx.resource_owner_id,
        transformer=NoOpTransformer(),
        description="Data owner requires no masking",
    ),
    # Rule 2: Customer service sees partial masking
    ConditionalRule(
        condition=lambda ctx: "customer_service" in ctx.subject.roles,
        transformer=PartialMaskingTransformer(prefix_len=3, suffix_len=4),
        description="Customer service sees partial phone",
    ),
    # Rule 3: Analysts see hash values
    ConditionalRule(
        condition=lambda ctx: "analyst" in ctx.subject.roles,
        transformer=HashReplacementTransformer(secret_key=ANALYST_KEY),
        description="Analysts see hash substitutes",
    ),
    # Rule 4: External API users get full masking outside business hours
    ConditionalRule(
        condition=lambda ctx: (
            ctx.environment.client_type == "external_api"
            and not ctx.environment.is_business_hours
        ),
        transformer=FullMaskingTransformer(),
        description="External API full masking outside business hours",
    ),
])

class MaskingPolicy(BaseModel):
    """Masking Policy Definition"""

    id: str = Field(description="Policy unique identifier")
    name: str = Field(description="Policy name")
    description: str = Field(default="", description="Policy description")

    target: MaskingTarget = Field(description="Masking target specification")
    mode: MaskingMode = Field(description="Masking mode")
    mode_config: dict = Field(default_factory=dict, description="Mode configuration")

    conditions: list[PolicyCondition] = Field(
        default_factory=list,
        description="Application conditions (AND relationship)",
    )

    priority: int = Field(default=100, description="Policy priority (lower = higher)")
    effective_from: datetime | None = Field(default=None)
    effective_until: datetime | None = Field(default=None)
    enabled: bool = Field(default=True)


class MaskingTarget(BaseModel):
    """Masking target specification"""

    object_types: list[str] = Field(
        default_factory=lambda: ["*"],
        description="Applicable Ontology object types",
    )
    field_patterns: list[str] = Field(
        description="Field name matching patterns (supports wildcards)",
    )
    classification_levels: list[str] = Field(
        default_factory=list,
        description="Applicable data classification levels",
    )

class MaskingPolicyRegistry:
    """Masking Policy Registry"""

    def __init__(self):
        self._policies: dict[str, MaskingPolicy] = {}
        self._compiled_cache: dict[str, CompiledPolicy] = {}

    def register(self, policy: MaskingPolicy) -> None:
        self._policies[policy.id] = policy
        self._invalidate_cache()

    def resolve_for_field(
        self,
        object_type: str,
        field_name: str,
        classification: str,
        context: RequestContext,
    ) -> MaskingTransformer | None:
        """Resolve applicable masking transformer for a field"""
        applicable = []
        for policy in self._policies.values():
            if not policy.enabled:
                continue
            if not self._matches_target(policy.target, object_type, field_name, classification):
                continue
            if not self._matches_conditions(policy.conditions, context):
                continue
            if not self._is_effective(policy):
                continue
            applicable.append(policy)

        if not applicable:
            return None

        applicable.sort(key=lambda p: p.priority)
        best = applicable[0]
        return self._create_transformer(best.mode, best.mode_config)

class CachedMaskingEngine(DynamicMaskingEngine):
    """Masking engine with caching"""

    async def apply_masking(
        self,
        result_set: ResultSet,
        request_context: RequestContext,
    ) -> ResultSet:
        cache_key = self._build_cache_key(request_context)

        field_masks = self._field_mask_cache.get(cache_key)
        if field_masks is None:
            field_masks = await self._resolve_field_masks(result_set, request_context)
            self._field_mask_cache.set(cache_key, field_masks, ttl=300)

        return self._batch_transform(result_set, field_masks)

class VectorizedMaskingEngine:
    """Vectorized masking engine using Arrow/Pandas for batch processing"""

    def transform_column(
        self,
        column: pa.Array,
        transformer: MaskingTransformer,
        context: RequestContext,
    ) -> pa.Array:
        """Vectorized column-level transformation"""
        if transformer.mode == MaskingMode.FULL:
            return pa.array(["***"] * len(column), type=pa.string())

        elif transformer.mode == MaskingMode.PARTIAL:
            values = column.to_pylist()
            masked = [transformer.transform(v, "", context) for v in values]
            return pa.array(masked, type=pa.string())

        elif transformer.mode == MaskingMode.HASH:
            values = column.to_pylist()
            masked = [transformer.transform(v, "", context) for v in values]
            return pa.array(masked, type=pa.string())

        return column

class ABACDrivenMaskingResolver:
    """ABAC-driven masking policy resolver"""

    CLEARANCE_MASKING_MAP = {
        # User clearance → Data classification → Masking mode
        ("PUBLIC", "A1"): MaskingMode.NONE,
        ("PUBLIC", "A2"): MaskingMode.PARTIAL,
        ("PUBLIC", "B"): MaskingMode.FULL,
        ("PUBLIC", "C"): MaskingMode.FULL,
        ("PUBLIC", "D"): MaskingMode.FULL,

        ("INTERNAL", "A1"): MaskingMode.NONE,
        ("INTERNAL", "A2"): MaskingMode.NONE,
        ("INTERNAL", "B"): MaskingMode.PARTIAL,
        ("INTERNAL", "C"): MaskingMode.HASH,
        ("INTERNAL", "D"): MaskingMode.FULL,

        ("CONFIDENTIAL", "A1"): MaskingMode.NONE,
        ("CONFIDENTIAL", "A2"): MaskingMode.NONE,
        ("CONFIDENTIAL", "B"): MaskingMode.NONE,
        ("CONFIDENTIAL", "C"): MaskingMode.PARTIAL,
        ("CONFIDENTIAL", "D"): MaskingMode.HASH,

        ("SECRET", "A1"): MaskingMode.NONE,
        ("SECRET", "A2"): MaskingMode.NONE,
        ("SECRET", "B"): MaskingMode.NONE,
        ("SECRET", "C"): MaskingMode.NONE,
        ("SECRET", "D"): MaskingMode.PARTIAL,
    }

class MaskingAuditEvent(BaseModel):
    """Masking audit event"""

    event_type: str = "DATA_MASKING"
    timestamp: datetime
    subject_id: str
    resource_id: str
    object_type: str
    fields_masked: list[FieldMaskingDetail]
    masking_policy_ids: list[str]
    query_id: str
    result_row_count: int

class FieldMaskingDetail(BaseModel):
    field_name: str
    masking_mode: MaskingMode
    classification_level: str
    original_type: str

class TestDynamicMasking:
    """Masking correctness tests"""

    def test_full_masking_hides_all_content(self):
        transformer = FullMaskingTransformer()
        assert transformer.transform("sensitive_data", "field", ctx) == "******"
        assert transformer.transform(12345, "field", ctx) == 0

    def test_partial_masking_preserves_structure(self):
        transformer = PartialMaskingTransformer(prefix_len=3, suffix_len=4)
        result = transformer.transform("13812345678", "phone", ctx)
        assert result == "138****5678"
        assert len(result) == 11

    def test_hash_deterministic(self):
        transformer = HashReplacementTransformer(secret_key=b"test")
        r1 = transformer.transform("value", "field", ctx)
        r2 = transformer.transform("value", "field", ctx)
        assert r1 == r2

    def test_hash_different_fields_different_output(self):
        transformer = HashReplacementTransformer(secret_key=b"test")
        r1 = transformer.transform("value", "field_a", ctx)
        r2 = transformer.transform("value", "field_b", ctx)
        assert r1 != r2

    def test_fpe_preserves_format(self):
        transformer = FPETransformer(key=b"0" * 16, tweak=b"0" * 8)
        result = transformer.transform("555-123-4567", "phone", ctx)
        assert len(result) == 12
        assert result[3] == "-"
        assert result[7] == "-"

    def test_conditional_masking_role_based(self):
        owner_ctx = make_context(subject_id="owner_1", resource_owner="owner_1")
        analyst_ctx = make_context(roles=["analyst"])

        result_owner = phone_masking.transform("13812345678", "phone", owner_ctx)
        result_analyst = phone_masking.transform("13812345678", "phone", analyst_ctx)

        assert result_owner == "13812345678"
        assert len(result_analyst) == 16

@pytest.mark.benchmark
def test_masking_performance_under_sla(benchmark):
    """Ensure masking latency stays within SLA"""
    engine = DynamicMaskingEngine(...)
    result_set = generate_result_set(rows=10000, cols=20, masked_cols=5)

    result = benchmark(engine.apply_masking, result_set, test_context)

    assert benchmark.stats["mean"] < 0.05  # 50ms for 10K rows