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

from typing import Any

from .schema import PolicyDefinition

# Array categories: first-match-wins rule lists (concatenated in order)
_ARRAY_CATEGORIES = (
    "file", "network", "commands", "env", "dns", "connect",
    "package_rules", "signal_rules", "unix_socket_rules",
)

# Object categories: plain objects (last-write-wins shallow merge)
_OBJECT_CATEGORIES = ("env_policy", "resource_limits", "audit_settings")


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:
    merged_data = base.model_dump(exclude_none=True)

    for override in overrides:
        # Array categories: concatenate in order
        for key in _ARRAY_CATEGORIES:
            override_rules = getattr(override, key, None)
            if override_rules:
                base_rules = merged_data.get(key) or []
                if mode == "append":
                    merged_data[key] = base_rules + list(override_rules)
                else:
                    merged_data[key] = list(override_rules) + base_rules

        # Object categories: shallow merge (override wins)
        for key in _OBJECT_CATEGORIES:
            override_obj = getattr(override, key, None)
            if override_obj is not None:
                existing: dict[str, Any] = merged_data.get(key) or {}
                if hasattr(override_obj, "model_dump"):
                    override_dict = override_obj.model_dump(exclude_none=True)
                else:
                    override_dict = dict(override_obj)
                merged_data[key] = {**existing, **override_dict}

    return PolicyDefinition.model_validate(merged_data)
