"""
Policy Evaluator
================

Engine for evaluating memory operations against policies.

The evaluator:
- Matches operations to applicable policies
- Evaluates policy rules
- Returns allow/deny decisions with explanations
- Logs policy decisions for audit
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any, Callable, Dict, List, Optional

from hieromem_core.policy.schema import (
    MemoryPolicy,
    RetentionPolicy,
)

logger = logging.getLogger(__name__)


class PolicyAction(Enum):
    """Actions that policies can evaluate."""
    INSERT = "insert"
    RETRIEVE = "retrieve"
    DELETE = "delete"
    UPDATE = "update"
    SUMMARIZE = "summarize"
    DECAY = "decay"
    PRUNE = "prune"


class PolicyResult(Enum):
    """Result of policy evaluation."""
    ALLOW = "allow"
    DENY = "deny"
    WARN = "warn"  # Allow but log warning


@dataclass
class PolicyDecision:
    """
    Result of evaluating an operation against policies.

    Contains the decision, reasoning, and any modifications
    that should be applied to the operation.
    """
    result: PolicyResult
    policy_name: str
    action: PolicyAction
    reason: str
    timestamp: datetime = field(default_factory=datetime.utcnow)

    # Modifications to apply (e.g., adjusted lambda)
    modifications: Dict[str, Any] = field(default_factory=dict)

    # Warnings to surface (even if allowed)
    warnings: List[str] = field(default_factory=list)

    # Audit metadata
    namespace: Optional[str] = None
    episode_id: Optional[str] = None

    def is_allowed(self) -> bool:
        """Check if the operation is allowed."""
        return self.result in (PolicyResult.ALLOW, PolicyResult.WARN)

    @property
    def allowed(self) -> bool:
        """Alias for is_allowed() for backwards compatibility."""
        return self.is_allowed()

    @property
    def reasons(self) -> List[str]:
        """Return reasons as a list for backwards compatibility."""
        return [self.reason] if self.reason else []

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for logging/serialization."""
        return {
            "result": self.result.value,
            "policy_name": self.policy_name,
            "action": self.action.value,
            "reason": self.reason,
            "timestamp": self.timestamp.isoformat(),
            "modifications": self.modifications,
            "warnings": self.warnings,
            "namespace": self.namespace,
            "episode_id": self.episode_id,
        }


@dataclass
class EvaluationContext:
    """
    Context for policy evaluation.

    Contains all information needed to evaluate an operation.
    """
    action: PolicyAction
    namespace: str

    # Episode metadata
    episode_id: Optional[str] = None
    content: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)

    # Trust/decay context
    current_trust: Optional[float] = None
    last_accessed: Optional[datetime] = None
    cycles_unused: int = 0

    # Graph context
    relations: List[str] = field(default_factory=list)
    edge_count: int = 0

    # Memory state
    total_episodes: int = 0


class PolicyEvaluator:
    """
    Evaluates a single policy against an operation.

    Use PolicyEngine for multi-policy evaluation with precedence.
    """

    def __init__(self, policy: MemoryPolicy):
        """
        Initialize evaluator with a policy.

        Args:
            policy: The MemoryPolicy to evaluate against.
        """
        self.policy = policy

    def evaluate(self, ctx: EvaluationContext) -> PolicyDecision:
        """
        Evaluate an operation against this policy.

        Args:
            ctx: Evaluation context with operation details.

        Returns:
            PolicyDecision with result and reasoning.
        """
        # Check if policy is enabled
        if not self.policy.enabled:
            return PolicyDecision(
                result=PolicyResult.ALLOW,
                policy_name=self.policy.name,
                action=ctx.action,
                reason="Policy disabled",
                namespace=ctx.namespace,
                episode_id=ctx.episode_id,
            )

        # Check namespace match
        if not self._matches_namespace(ctx.namespace):
            return PolicyDecision(
                result=PolicyResult.ALLOW,
                policy_name=self.policy.name,
                action=ctx.action,
                reason="Namespace not matched",
                namespace=ctx.namespace,
                episode_id=ctx.episode_id,
            )

        # Dispatch to action-specific evaluation
        evaluators = {
            PolicyAction.INSERT: self._evaluate_insert,
            PolicyAction.RETRIEVE: self._evaluate_retrieve,
            PolicyAction.DELETE: self._evaluate_delete,
            PolicyAction.UPDATE: self._evaluate_update,
            PolicyAction.SUMMARIZE: self._evaluate_summarize,
            PolicyAction.DECAY: self._evaluate_decay,
            PolicyAction.PRUNE: self._evaluate_prune,
        }

        evaluator = evaluators.get(ctx.action)
        if evaluator:
            return evaluator(ctx)

        # Default allow for unknown actions
        return PolicyDecision(
            result=PolicyResult.ALLOW,
            policy_name=self.policy.name,
            action=ctx.action,
            reason="Unknown action type",
            namespace=ctx.namespace,
            episode_id=ctx.episode_id,
        )

    def _matches_namespace(self, namespace: str) -> bool:
        """Check if policy applies to namespace."""
        if self.policy.namespace == "*":
            return True
        return self.policy.namespace == namespace

    def _evaluate_insert(self, ctx: EvaluationContext) -> PolicyDecision:
        """Evaluate an insert operation."""
        warnings = []

        # Check governance: forbidden relations
        gov = self.policy.governance
        for rel in ctx.relations:
            if gov.is_relation_forbidden(rel):
                return PolicyDecision(
                    result=PolicyResult.DENY,
                    policy_name=self.policy.name,
                    action=ctx.action,
                    reason=f"Forbidden relation: {rel}",
                    namespace=ctx.namespace,
                    episode_id=ctx.episode_id,
                )

        # Check governance: blocked content
        if ctx.content and gov.is_content_blocked(ctx.content):
            return PolicyDecision(
                result=PolicyResult.DENY,
                policy_name=self.policy.name,
                action=ctx.action,
                reason="Content matches blocked pattern",
                namespace=ctx.namespace,
                episode_id=ctx.episode_id,
            )

        # Check governance: capacity limit
        if gov.max_episodes is not None and ctx.total_episodes >= gov.max_episodes:
            return PolicyDecision(
                result=PolicyResult.DENY,
                policy_name=self.policy.name,
                action=ctx.action,
                reason=f"Episode limit exceeded: {ctx.total_episodes} >= {gov.max_episodes}",
                namespace=ctx.namespace,
                episode_id=ctx.episode_id,
            )

        # Warn if approaching capacity
        if gov.max_episodes is not None:
            usage = ctx.total_episodes / gov.max_episodes
            if usage > 0.9:
                warnings.append(f"Approaching capacity: {usage:.0%} used")

        return PolicyDecision(
            result=PolicyResult.WARN if warnings else PolicyResult.ALLOW,
            policy_name=self.policy.name,
            action=ctx.action,
            reason="Insert allowed" if not warnings else "Insert allowed with warnings",
            warnings=warnings,
            namespace=ctx.namespace,
            episode_id=ctx.episode_id,
        )

    def _evaluate_retrieve(self, ctx: EvaluationContext) -> PolicyDecision:
        """Evaluate a retrieve operation."""
        # Retrieves are generally allowed; compute trust boost
        modifications = {}

        # Apply access boost (bounded by governance)
        boost = min(
            self.policy.retention.access_boost,
            self.policy.governance.max_trust_boost
        )
        modifications["trust_boost"] = boost

        return PolicyDecision(
            result=PolicyResult.ALLOW,
            policy_name=self.policy.name,
            action=ctx.action,
            reason="Retrieve allowed",
            modifications=modifications,
            namespace=ctx.namespace,
            episode_id=ctx.episode_id,
        )

    def _evaluate_delete(self, ctx: EvaluationContext) -> PolicyDecision:
        """Evaluate a delete operation."""
        # Check if audit is required
        if self.policy.governance.audit_required:
            return PolicyDecision(
                result=PolicyResult.WARN,
                policy_name=self.policy.name,
                action=ctx.action,
                reason="Delete allowed (audit logged)",
                warnings=["Audit: Delete operation logged"],
                namespace=ctx.namespace,
                episode_id=ctx.episode_id,
            )

        return PolicyDecision(
            result=PolicyResult.ALLOW,
            policy_name=self.policy.name,
            action=ctx.action,
            reason="Delete allowed",
            namespace=ctx.namespace,
            episode_id=ctx.episode_id,
        )

    def _evaluate_update(self, ctx: EvaluationContext) -> PolicyDecision:
        """Evaluate an update operation."""
        # Same checks as insert
        return self._evaluate_insert(ctx)

    def _evaluate_summarize(self, ctx: EvaluationContext) -> PolicyDecision:
        """Evaluate a summarization operation."""
        summ = self.policy.summarization

        # Check if summarization should trigger
        should_summarize = False
        reason_parts = []

        if summ.max_episodes and ctx.total_episodes > summ.max_episodes:
            should_summarize = True
            reason_parts.append(f"episode count ({ctx.total_episodes}) exceeds max ({summ.max_episodes})")

        modifications = {
            "depth": summ.depth,
            "preserve_graph": summ.preserve_graph,
            "min_fidelity": summ.min_fidelity,
            "batch_size": summ.batch_size,
        }

        reason_str = (
            f"Summarization triggered: {', '.join(reason_parts)}"
            if should_summarize else "Summarization not required"
        )
        return PolicyDecision(
            result=PolicyResult.ALLOW,
            policy_name=self.policy.name,
            action=ctx.action,
            reason=reason_str,
            modifications=modifications,
            namespace=ctx.namespace,
            episode_id=ctx.episode_id,
        )

    def _evaluate_decay(self, ctx: EvaluationContext) -> PolicyDecision:
        """Evaluate a decay operation and compute new trust."""
        ret = self.policy.retention
        priority = self.policy.priority
        gov = self.policy.governance

        # Get base lambda
        base_lambda = ret.base_lambda

        # Apply priority multiplier
        multiplier = priority.get_priority_multiplier(ctx.metadata)
        if multiplier == float("inf"):
            # Critical memory - no decay
            return PolicyDecision(
                result=PolicyResult.ALLOW,
                policy_name=self.policy.name,
                action=ctx.action,
                reason="Critical memory - no decay",
                modifications={"new_trust": ctx.current_trust or 1.0, "decayed": False},
                namespace=ctx.namespace,
                episode_id=ctx.episode_id,
            )

        # Compute effective lambda (bounded by governance)
        effective_lambda = min(
            max(base_lambda * multiplier, gov.min_lambda),
            gov.max_lambda
        )

        # Compute new trust
        current = ctx.current_trust or 1.0
        new_trust = ret.compute_decay(current, ctx.cycles_unused)

        # Adjust for effective lambda (recompute if different from base)
        if effective_lambda != ret.base_lambda:
            # Create temporary policy with adjusted lambda
            temp_ret = RetentionPolicy(
                decay_curve=ret.decay_curve,
                base_lambda=effective_lambda,
                min_trust=ret.min_trust,
            )
            new_trust = temp_ret.compute_decay(current, ctx.cycles_unused)

        # Check if should prune
        should_prune = new_trust < ret.min_trust

        return PolicyDecision(
            result=PolicyResult.ALLOW,
            policy_name=self.policy.name,
            action=ctx.action,
            reason=f"Decay: {current:.3f} -> {new_trust:.3f} (λ={effective_lambda:.3f})",
            modifications={
                "new_trust": new_trust,
                "effective_lambda": effective_lambda,
                "should_prune": should_prune,
                "decayed": True,
            },
            namespace=ctx.namespace,
            episode_id=ctx.episode_id,
        )

    def _evaluate_prune(self, ctx: EvaluationContext) -> PolicyDecision:
        """Evaluate a prune operation."""
        graph = self.policy.graph
        warnings = []

        # Check protected edge types
        for rel in ctx.relations:
            if rel in graph.protected_edge_types:
                warnings.append(f"Protected edge type present: {rel}")

        # Check orphan pruning setting
        if not graph.prune_orphans and ctx.edge_count == 0:
            return PolicyDecision(
                result=PolicyResult.DENY,
                policy_name=self.policy.name,
                action=ctx.action,
                reason="Orphan pruning disabled",
                namespace=ctx.namespace,
                episode_id=ctx.episode_id,
            )

        return PolicyDecision(
            result=PolicyResult.WARN if warnings else PolicyResult.ALLOW,
            policy_name=self.policy.name,
            action=ctx.action,
            reason="Prune allowed" if not warnings else "Prune allowed with warnings",
            warnings=warnings,
            namespace=ctx.namespace,
            episode_id=ctx.episode_id,
        )


class PolicyEngine:
    """
    Engine for evaluating operations against multiple policies.

    Policies are evaluated in order of specificity:
    1. Namespace-specific policies (exact match)
    2. Wildcard policies (namespace="*")

    First deny wins; if all allow, operation proceeds.
    """

    def __init__(self) -> None:
        """Initialize the policy engine."""
        self._policies: Dict[str, MemoryPolicy] = {}
        self._decision_handlers: List[Callable[[PolicyDecision], None]] = []

    def register_policy(self, policy: MemoryPolicy) -> None:
        """
        Register a policy with the engine.

        Args:
            policy: The MemoryPolicy to register.
        """
        self._policies[policy.name] = policy
        logger.info(f"Registered policy: {policy.name} (namespace={policy.namespace})")

    def unregister_policy(self, name: str) -> bool:
        """
        Unregister a policy by name.

        Args:
            name: Policy name to unregister.

        Returns:
            True if policy was removed, False if not found.
        """
        if name in self._policies:
            del self._policies[name]
            logger.info(f"Unregistered policy: {name}")
            return True
        return False

    def get_policy(self, name: str) -> Optional[MemoryPolicy]:
        """Get a policy by name."""
        return self._policies.get(name)

    def list_policies(self) -> List[MemoryPolicy]:
        """List all registered policies."""
        return list(self._policies.values())

    def add_decision_handler(self, handler: Callable[[PolicyDecision], None]) -> None:
        """
        Add a handler to be called for every policy decision.

        Useful for logging, metrics, or audit trails.

        Args:
            handler: Callable that receives PolicyDecision objects.
        """
        self._decision_handlers.append(handler)

    def evaluate(self, ctx: EvaluationContext) -> PolicyDecision:
        """
        Evaluate an operation against all applicable policies.

        Args:
            ctx: Evaluation context.

        Returns:
            Final PolicyDecision (first deny or merged allow).
        """
        # Get applicable policies (sorted by specificity)
        applicable = self._get_applicable_policies(ctx.namespace)

        if not applicable:
            # No policies - default allow
            decision = PolicyDecision(
                result=PolicyResult.ALLOW,
                policy_name="<none>",
                action=ctx.action,
                reason="No applicable policies",
                namespace=ctx.namespace,
                episode_id=ctx.episode_id,
            )
            self._notify_handlers(decision)
            return decision

        # Evaluate each policy
        all_decisions: List[PolicyDecision] = []
        merged_modifications: Dict[str, Any] = {}
        all_warnings: List[str] = []

        for policy in applicable:
            evaluator = PolicyEvaluator(policy)
            decision = evaluator.evaluate(ctx)
            all_decisions.append(decision)

            # Log decision
            self._notify_handlers(decision)

            # First deny wins
            if decision.result == PolicyResult.DENY:
                return decision

            # Collect modifications and warnings
            merged_modifications.update(decision.modifications)
            all_warnings.extend(decision.warnings)

        # All allowed - return merged result
        return PolicyDecision(
            result=PolicyResult.WARN if all_warnings else PolicyResult.ALLOW,
            policy_name=",".join(p.name for p in applicable),
            action=ctx.action,
            reason="Allowed by all policies",
            modifications=merged_modifications,
            warnings=all_warnings,
            namespace=ctx.namespace,
            episode_id=ctx.episode_id,
        )

    def _get_applicable_policies(self, namespace: str) -> List[MemoryPolicy]:
        """
        Get policies applicable to a namespace, sorted by specificity.

        Namespace-specific policies come before wildcards.
        """
        specific = []
        wildcards = []

        for policy in self._policies.values():
            if not policy.enabled:
                continue
            if policy.namespace == namespace:
                specific.append(policy)
            elif policy.namespace == "*":
                wildcards.append(policy)

        return specific + wildcards

    def _notify_handlers(self, decision: PolicyDecision) -> None:
        """Notify all registered handlers of a decision."""
        for handler in self._decision_handlers:
            try:
                handler(decision)
            except Exception as e:
                logger.warning(f"Decision handler error: {e}")

    def check_insert(
        self,
        namespace: str,
        content: Optional[str] = None,
        relations: Optional[List[str]] = None,
        total_episodes: int = 0,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> PolicyDecision:
        """
        Convenience method for checking insert operations.

        Args:
            namespace: Target namespace.
            content: Content being inserted.
            relations: Graph relations being added.
            total_episodes: Current episode count.
            metadata: Episode metadata.

        Returns:
            PolicyDecision for the insert operation.
        """
        ctx = EvaluationContext(
            action=PolicyAction.INSERT,
            namespace=namespace,
            content=content,
            relations=relations or [],
            total_episodes=total_episodes,
            metadata=metadata or {},
        )
        return self.evaluate(ctx)

    def check_decay(
        self,
        namespace: str,
        episode_id: str,
        current_trust: float,
        cycles_unused: int = 1,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> PolicyDecision:
        """
        Convenience method for checking decay operations.

        Args:
            namespace: Episode namespace.
            episode_id: Episode identifier.
            current_trust: Current trust score.
            cycles_unused: Number of decay cycles since last access.
            metadata: Episode metadata.

        Returns:
            PolicyDecision with new_trust in modifications.
        """
        ctx = EvaluationContext(
            action=PolicyAction.DECAY,
            namespace=namespace,
            episode_id=episode_id,
            current_trust=current_trust,
            cycles_unused=cycles_unused,
            metadata=metadata or {},
        )
        return self.evaluate(ctx)


class MemoryPolicyEvaluator:
    """
    Memory-specific policy evaluator.

    Provides convenience methods for common memory operations
    (insert, summarize, forget) that wrap the PolicyEngine.
    """

    def __init__(
        self,
        policy: MemoryPolicy,
        event_log: Optional[Any] = None,
    ) -> None:
        """
        Initialize the memory policy evaluator.

        Args:
            policy: The MemoryPolicy to evaluate against.
            event_log: Optional event log for audit (currently unused).
        """
        self._policy = policy
        self._event_log = event_log
        self._engine = PolicyEngine()
        self._engine.register_policy(policy)

    def evaluate_insert(
        self,
        graph: Any = None,
        total_episodes: int = 0,
        namespace: Optional[str] = None,
        vector: Any = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> PolicyDecision:
        """
        Evaluate an insert operation against the policy.

        Args:
            graph: The knowledge graph being inserted.
            total_episodes: Current total episodes in memory.
            namespace: Namespace for the operation.
            vector: The embedding vector being inserted.
            metadata: Optional episode metadata.

        Returns:
            PolicyDecision indicating whether insert is allowed.
        """
        # Extract relations from graph if present
        relations: List[str] = []
        if graph is not None:
            try:
                import networkx as nx
                if isinstance(graph, nx.DiGraph) or isinstance(graph, nx.Graph):
                    for _, _, data in graph.edges(data=True):
                        if "relation" in data:
                            relations.append(data["relation"])
            except Exception:
                pass

        # Check for forbidden relations
        if self._policy.governance and self._policy.governance.forbidden_relations:
            forbidden = set(self._policy.governance.forbidden_relations)
            found_forbidden = forbidden.intersection(set(relations))
            if found_forbidden:
                return PolicyDecision(
                    result=PolicyResult.DENY,
                    policy_name=self._policy.name,
                    action=PolicyAction.INSERT,
                    reason=f"Graph contains forbidden relations: {found_forbidden}",
                    namespace=namespace or self._policy.namespace,
                )

        # Check max episodes
        if self._policy.governance and self._policy.governance.max_episodes:
            if total_episodes >= self._policy.governance.max_episodes:
                return PolicyDecision(
                    result=PolicyResult.DENY,
                    policy_name=self._policy.name,
                    action=PolicyAction.INSERT,
                    reason=f"Maximum episodes ({self._policy.governance.max_episodes}) reached",
                    namespace=namespace or self._policy.namespace,
                )

        ctx = EvaluationContext(
            action=PolicyAction.INSERT,
            namespace=namespace or self._policy.namespace,
            metadata=metadata or {},
            relations=relations,
            total_episodes=total_episodes,
        )
        return self._engine.evaluate(ctx)

    def should_summarise(self, episodes: List[Any]) -> bool:
        """
        Check if summarization should be triggered.

        Args:
            episodes: Current list of episodes in memory.

        Returns:
            True if summarization should be triggered.
        """
        if self._policy.summarization is None:
            return False

        # Check episode count trigger using max_episodes
        if self._policy.summarization.max_episodes is not None:
            if len(episodes) >= self._policy.summarization.max_episodes:
                return True

        return False

    def retention_candidates(self, episodes: List[Any]) -> List[Any]:
        """
        Identify episodes that are candidates for retention/pruning.

        Evaluates each episode against the retention policy and returns
        those that should be considered for removal (low trust, old age, etc.)

        Args:
            episodes: List of episodes to evaluate.

        Returns:
            List of episodes that are candidates for pruning.
        """
        if not episodes:
            return []

        candidates = []
        retention = self._policy.retention

        for episode in episodes:
            should_prune = False
            trust = getattr(episode, 'trust', 1.0)

            # Check trust threshold
            if retention and retention.min_trust is not None:
                if trust < retention.min_trust:
                    should_prune = True

            # Check episode age against horizon
            if retention and retention.horizon is not None:
                tocc = getattr(episode, 'tocc', None)
                if tocc is not None:
                    from datetime import datetime
                    age = datetime.utcnow() - tocc
                    if age > retention.horizon:
                        should_prune = True

            if should_prune:
                candidates.append(episode)

        return candidates

    def evaluate_forget(
        self,
        episode_id: Optional[str] = None,
        episode_ids: Optional[List[str]] = None,
        reason: str = "manual",
        namespace: Optional[str] = None,
    ) -> PolicyDecision:
        """
        Evaluate a forget/delete operation against the policy.

        Args:
            episode_id: ID of the episode being forgotten (singular).
            episode_ids: IDs of episodes being forgotten (plural).
            reason: Reason for forgetting.
            namespace: Optional namespace override.

        Returns:
            PolicyDecision indicating whether forget is allowed.
        """
        # Support both singular and plural forms
        ids = episode_ids or ([episode_id] if episode_id else [])

        ctx = EvaluationContext(
            action=PolicyAction.DELETE,
            namespace=namespace or self._policy.namespace,
            episode_id=ids[0] if ids else None,
            metadata={"reason": reason, "episode_ids": ids},
        )
        return self._engine.evaluate(ctx)

    def evaluate_summarise(
        self,
        episodes: List[Any],
        namespace: Optional[str] = None,
    ) -> PolicyDecision:
        """
        Evaluate a summarization operation against the policy.

        Args:
            episodes: Episodes being summarized.
            namespace: Optional namespace override.

        Returns:
            PolicyDecision indicating whether summarization is allowed.
        """
        ctx = EvaluationContext(
            action=PolicyAction.SUMMARIZE,
            namespace=namespace or self._policy.namespace,
            metadata={"episode_count": len(episodes)},
        )
        return self._engine.evaluate(ctx)
