# src/agentsh_secure_sandbox/policies/merge.py
from __future__ import annotations

from typing import Any

from .schema import PolicyDefinition

_CATEGORIES = ("file", "network", "commands", "env", "dns", "connect")


def merge(base: PolicyDefinition, *overrides: PolicyDefinition) -> PolicyDefinition:
    """Merge policy extensions AFTER base rules (base takes priority in first-match-wins)."""
    return _merge_internal(base, overrides, "append")


def merge_prepend(base: PolicyDefinition, *overrides: PolicyDefinition) -> PolicyDefinition:
    """Merge overrides BEFORE base rules (overrides take priority in first-match-wins)."""
    return _merge_internal(base, overrides, "prepend")


def _merge_internal(
    base: PolicyDefinition,
    overrides: tuple[PolicyDefinition, ...],
    mode: str,
) -> PolicyDefinition:
    updates: dict[str, list[Any]] = {}
    # For prepend mode, iterate in reverse so that earlier overrides
    # end up with higher priority (closer to the front).
    ordered = reversed(overrides) if mode == "prepend" else overrides
    for override in ordered:
        for key in _CATEGORIES:
            override_rules = getattr(override, key, None)
            if override_rules:
                base_rules = updates.get(key) or getattr(base, key, None) or []
                if mode == "append":
                    updates[key] = base_rules + override_rules
                else:
                    updates[key] = override_rules + base_rules
    merged_data = base.model_dump(exclude_none=True)
    merged_data.update(updates)
    return PolicyDefinition.model_validate(merged_data)
