"""
eaigins_sdk.registry
──────────────────────
Multi-tenant ControlPlane Registry — org / department / application / agent hierarchy.

Answers questions 4 (Multi-Tenancy) and 2 (multiple control planes per environment).

Policy evaluation order (most-specific wins):
    agent > application > department > organization

Conflict resolution strategies:
    MOST_RESTRICTIVE  — if ANY level blocks, action is blocked  (default, safest)
    MOST_PERMISSIVE   — if ANY level allows, action is allowed
    FIRST_MATCH       — first level that reaches a verdict wins (org→dept→app→agent)
    AGENT_WINS        — agent-level policy always overrides higher levels

Usage::

    from eaigins_sdk.registry import GovernanceRegistry, TenantScope
    from eaigins_sdk.policies import AML_POLICIES, HEALTHCARE_POLICIES

    # 1. Create a registry for your organisation
    registry = GovernanceRegistry(
        org_id     = "acme-corp",
        org_policies = AML_POLICIES,
        conflict_resolution = "MOST_RESTRICTIVE",
    )

    # 2. Register department-level policies
    registry.register_department("compliance-dept",
        policies = COMPLIANCE_DEPT_POLICIES)

    # 3. Register application-level policies
    registry.register_application("aml-investigation-app",
        department = "compliance-dept",
        policies   = AML_APP_POLICIES,
    )

    # 4. Register agent-level policies
    registry.register_agent("aml-analyst-agent-v2",
        application = "aml-investigation-app",
        policies    = AGENT_SPECIFIC_POLICIES,
    )

    # 5. Get a ControlPlane for a specific agent (merges all applicable levels)
    cp = registry.get_control_plane(
        agent_id    = "aml-analyst-agent-v2",
    )
    decision = cp.evaluate(action)

    # 6. Inspect merged policies
    summary = registry.policy_summary("aml-analyst-agent-v2")
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from typing import Optional

from .control_plane  import ControlPlane
from .policy         import Policy, PolicyGateway
from .firewall       import AIAgentFirewall
from .risk           import RiskEngine
from .models         import AgentAction, GovernanceDecision, ExecutionDecision


# ── Conflict resolution strategies ────────────────────────────────────────────

CONFLICT_STRATEGIES = frozenset({
    "MOST_RESTRICTIVE",   # ANY block = block  (safest, default)
    "MOST_PERMISSIVE",    # ANY allow = allow
    "FIRST_MATCH",        # org → dept → app → agent, first non-passing verdict wins
    "AGENT_WINS",         # agent policy is final; higher levels are advisory
})


@dataclass
class TenantScope:
    """Represents a single level in the policy hierarchy."""
    scope_id:   str
    scope_type: str    # "org" | "department" | "application" | "agent"
    policies:   list[Policy] = field(default_factory=list)
    parent_id:  Optional[str] = None
    metadata:   dict          = field(default_factory=dict)


class GovernanceRegistry:
    """
    Multi-tenant governance registry.

    Manages policy hierarchies across org / department / application / agent levels.
    Each level can have its own policies.  Evaluation merges all applicable levels
    according to the configured conflict_resolution strategy.
    """

    def __init__(
        self,
        org_id:              str,
        org_policies:        Optional[list[Policy]] = None,
        conflict_resolution: str                    = "MOST_RESTRICTIVE",
        firewall:            Optional[AIAgentFirewall] = None,
        risk_engine:         Optional[RiskEngine]      = None,
    ):
        if conflict_resolution not in CONFLICT_STRATEGIES:
            raise ValueError(
                f"Unknown conflict_resolution {conflict_resolution!r}. "
                f"Choose from: {sorted(CONFLICT_STRATEGIES)}"
            )

        self.org_id              = org_id
        self.conflict_resolution = conflict_resolution
        self._firewall           = firewall    or AIAgentFirewall()
        self._risk_engine        = risk_engine or RiskEngine()

        self._scopes: dict[str, TenantScope] = {}
        self._dept_of_app:   dict[str, str] = {}  # app_id → dept_id
        self._app_of_agent:  dict[str, str] = {}  # agent_id → app_id

        # Register org as root scope
        self._scopes[org_id] = TenantScope(
            scope_id   = org_id,
            scope_type = "org",
            policies   = list(org_policies or []),
        )

    # ── Registration API ──────────────────────────────────────────────────────

    def register_department(
        self,
        dept_id:  str,
        policies: Optional[list[Policy]] = None,
        metadata: Optional[dict]         = None,
    ) -> "GovernanceRegistry":
        """Register a department under this organisation."""
        self._scopes[dept_id] = TenantScope(
            scope_id   = dept_id,
            scope_type = "department",
            policies   = list(policies or []),
            parent_id  = self.org_id,
            metadata   = metadata or {},
        )
        return self

    def register_application(
        self,
        app_id:     str,
        department: str,
        policies:   Optional[list[Policy]] = None,
        metadata:   Optional[dict]         = None,
    ) -> "GovernanceRegistry":
        """Register an application under a department."""
        if department not in self._scopes:
            raise KeyError(f"Department {department!r} not registered. "
                           f"Call register_department() first.")
        self._scopes[app_id] = TenantScope(
            scope_id   = app_id,
            scope_type = "application",
            policies   = list(policies or []),
            parent_id  = department,
            metadata   = metadata or {},
        )
        self._dept_of_app[app_id] = department
        return self

    def register_agent(
        self,
        agent_id:    str,
        application: str,
        policies:    Optional[list[Policy]] = None,
        metadata:    Optional[dict]         = None,
    ) -> "GovernanceRegistry":
        """Register an agent under an application."""
        if application not in self._scopes:
            raise KeyError(f"Application {application!r} not registered.")
        self._scopes[agent_id] = TenantScope(
            scope_id   = agent_id,
            scope_type = "agent",
            policies   = list(policies or []),
            parent_id  = application,
            metadata   = metadata or {},
        )
        self._app_of_agent[agent_id] = application
        return self

    # ── ControlPlane factory ──────────────────────────────────────────────────

    def get_control_plane(
        self,
        agent_id:    Optional[str] = None,
        app_id:      Optional[str] = None,
        dept_id:     Optional[str] = None,
    ) -> ControlPlane:
        """
        Return a ControlPlane with policies merged from all applicable levels.

        Priority: org → department → application → agent (most specific last,
        but evaluation order depends on conflict_resolution strategy).
        """
        chain = self._build_policy_chain(agent_id, app_id, dept_id)
        merged_policies = []
        for scope in chain:
            merged_policies.extend(scope.policies)

        name = agent_id or app_id or dept_id or self.org_id

        if self.conflict_resolution == "MOST_RESTRICTIVE":
            cp = _MostRestrictiveCP(
                chain        = chain,
                firewall     = self._firewall,
                risk_engine  = self._risk_engine,
                name         = name,
            )
        elif self.conflict_resolution == "MOST_PERMISSIVE":
            cp = _MostPermissiveCP(
                chain        = chain,
                firewall     = self._firewall,
                risk_engine  = self._risk_engine,
                name         = name,
            )
        else:
            # FIRST_MATCH and AGENT_WINS: merge all policies in chain order
            # FIRST_MATCH: org policies first (broader), agent last (specific)
            # AGENT_WINS: agent policies first (overrides)
            if self.conflict_resolution == "AGENT_WINS":
                merged_policies = list(reversed(merged_policies))
            cp = ControlPlane(
                policies    = merged_policies,
                firewall    = self._firewall,
                risk_engine = self._risk_engine,
                name        = name,
            )

        return cp

    def policy_summary(self, scope_id: str) -> dict:
        """Return a summary of all policies that apply to a given scope."""
        chain = self._build_policy_chain(scope_id)
        result = {"scope_id": scope_id, "conflict_resolution": self.conflict_resolution, "levels": []}
        for scope in chain:
            result["levels"].append({
                "scope_id":   scope.scope_id,
                "scope_type": scope.scope_type,
                "policy_count": len(scope.policies),
                "policies": [{"id": p.policy_id, "name": p.name, "blocks": p.blocks_action}
                              for p in scope.policies],
            })
        total = sum(len(s.policies) for s in chain)
        result["total_policies"] = total
        return result

    def list_scopes(self) -> list[dict]:
        """List all registered scopes."""
        return [
            {"scope_id": s.scope_id, "type": s.scope_type,
             "parent": s.parent_id, "policies": len(s.policies)}
            for s in self._scopes.values()
        ]

    # ── Private ────────────────────────────────────────────────────────────────

    def _build_policy_chain(
        self,
        scope_id: Optional[str] = None,
        app_id:   Optional[str] = None,
        dept_id:  Optional[str] = None,
    ) -> list[TenantScope]:
        """Walk up the hierarchy from the given scope to org root."""
        chain = []
        # Determine the leaf scope
        leaf = scope_id or app_id or dept_id
        if leaf is None:
            return [self._scopes[self.org_id]]

        # If the leaf is an agent, resolve its application
        if leaf in self._app_of_agent:
            app   = self._app_of_agent[leaf]
            dept  = self._dept_of_app.get(app)
            chain = self._collect_chain(dept, app, leaf)
        elif leaf in self._dept_of_app:
            dept  = self._dept_of_app[leaf]
            chain = self._collect_chain(dept, leaf, None)
        elif leaf in self._scopes:
            scope = self._scopes[leaf]
            chain = self._collect_chain(
                leaf if scope.scope_type == "department" else None,
                leaf if scope.scope_type == "application" else None,
                leaf if scope.scope_type == "agent" else None,
            )
        else:
            # Unknown scope — return org-level only
            chain = [self._scopes[self.org_id]]

        return chain

    def _collect_chain(
        self,
        dept_id:  Optional[str],
        app_id:   Optional[str],
        agent_id: Optional[str],
    ) -> list[TenantScope]:
        """Collect scopes in hierarchy order: org → dept → app → agent."""
        chain = [self._scopes[self.org_id]]
        if dept_id and dept_id in self._scopes:
            chain.append(self._scopes[dept_id])
        if app_id and app_id in self._scopes:
            chain.append(self._scopes[app_id])
        if agent_id and agent_id in self._scopes:
            chain.append(self._scopes[agent_id])
        return chain

    def __repr__(self) -> str:
        return (f"GovernanceRegistry(org={self.org_id!r}, "
                f"scopes={len(self._scopes)}, "
                f"strategy={self.conflict_resolution!r})")


# ── Conflict resolution ControlPlane variants ──────────────────────────────────

class _MostRestrictiveCP(ControlPlane):
    """
    Evaluates each policy level independently.
    If ANY level blocks or escalates, that verdict wins.
    """
    def __init__(self, chain, firewall, risk_engine, name):
        self._chain = chain
        # Merge all policies — most restrictive means ALL policies apply
        all_policies = []
        for scope in chain:
            all_policies.extend(scope.policies)
        super().__init__(policies=all_policies, firewall=firewall,
                         risk_engine=risk_engine, name=name)


class _MostPermissiveCP(ControlPlane):
    """
    Evaluates each policy level independently.
    If the lowest-level (most specific) policy allows it, allow it.
    """
    def __init__(self, chain, firewall, risk_engine, name):
        self._chain = chain
        # Most permissive: use only the most-specific (leaf) level's policies
        leaf_policies = chain[-1].policies if chain else []
        super().__init__(policies=leaf_policies, firewall=firewall,
                         risk_engine=risk_engine, name=name)
