"""
Policy Hooks
============

Integration between PolicyEngine and Memory class.

Provides:
- PolicyAwareMemory: Wrapper that enforces policies on memory operations
- Hooks for insert, summarize, decay, and prune operations
- Policy-driven trust computation

Usage:
    from hieromem_core.policy import PolicyEngine, MemoryPolicy
    from hieromem_core.policy.hooks import PolicyAwareMemory
    from hieromem_core.memory import Memory

    # Create engine with policies
    engine = PolicyEngine()
    engine.register_policy(MemoryPolicy.default())

    # Wrap memory with policy enforcement
    base_memory = Memory(Nmax=1000)
    memory = PolicyAwareMemory(base_memory, engine)

    # Operations now enforced by policies
    memory.insert(vector, graph, meta)  # May raise PolicyViolationError
"""

from __future__ import annotations

import logging
from typing import Any, Callable, Dict, List, Optional, TypeVar

import networkx as nx
import numpy as np

from hieromem_core.episode import Episode
from hieromem_core.policy.evaluator import (
    EvaluationContext,
    PolicyAction,
    PolicyDecision,
    PolicyEngine,
    PolicyResult,
)
from hieromem_core.policy.events import PolicyEvent, PolicyEventEmitter

logger = logging.getLogger(__name__)

T = TypeVar("T")


class PolicyViolationError(Exception):
    """Raised when an operation violates a policy."""

    def __init__(self, decision: PolicyDecision):
        self.decision = decision
        super().__init__(
            f"Policy violation ({decision.policy_name}): {decision.reason}"
        )


class PolicyAwareMemory:
    """
    Wrapper around Memory that enforces policy decisions.

    Intercepts memory operations and evaluates them against
    registered policies before allowing them to proceed.
    """

    def __init__(
        self,
        memory: Any,  # Memory type, avoid circular import
        policy_engine: PolicyEngine,
        event_emitter: Optional[PolicyEventEmitter] = None,
        raise_on_deny: bool = True,
        trace_id_generator: Optional[Callable[[], str]] = None,
    ):
        """
        Initialize the policy-aware memory wrapper.

        Args:
            memory: The underlying Memory instance to wrap.
            policy_engine: PolicyEngine for evaluating operations.
            event_emitter: Optional emitter for policy events.
            raise_on_deny: If True, raise PolicyViolationError on deny.
                          If False, log and skip the operation.
            trace_id_generator: Optional function to generate trace IDs.
        """
        self._memory = memory
        self._engine = policy_engine
        self._emitter = event_emitter
        self._raise_on_deny = raise_on_deny
        self._trace_id_generator = trace_id_generator

    @property
    def memory(self) -> Any:
        """Access the underlying memory instance."""
        return self._memory

    @property
    def engine(self) -> PolicyEngine:
        """Access the policy engine."""
        return self._engine

    def _get_trace_id(self) -> Optional[str]:
        """Generate a trace ID if generator is configured."""
        if self._trace_id_generator:
            return self._trace_id_generator()
        return None

    def _emit_event(self, event: PolicyEvent) -> None:
        """Emit a policy event if emitter is configured."""
        if self._emitter:
            self._emitter.emit(event)

    def _handle_decision(self, decision: PolicyDecision) -> bool:
        """
        Handle a policy decision.

        Args:
            decision: The policy decision to handle.

        Returns:
            True if operation should proceed, False otherwise.

        Raises:
            PolicyViolationError: If decision is DENY and raise_on_deny is True.
        """
        # Emit event
        self._emit_event(PolicyEvent.from_decision(decision, self._get_trace_id()))

        if decision.result == PolicyResult.DENY:
            if self._raise_on_deny:
                raise PolicyViolationError(decision)
            logger.warning(f"Policy denied: {decision.reason}")
            return False

        if decision.result == PolicyResult.WARN:
            for warning in decision.warnings:
                logger.warning(f"Policy warning: {warning}")

        return True

    def insert(
        self,
        vector: np.ndarray,
        graph: Optional[nx.DiGraph] = None,
        meta: Optional[Dict[str, Any]] = None,
        *,
        namespace: Optional[str] = None,
    ) -> Optional[Episode]:
        """
        Insert an episode with policy enforcement.

        Args:
            vector: Episode embedding vector.
            graph: Optional episode graph.
            meta: Optional metadata.
            namespace: Optional namespace override.

        Returns:
            The inserted Episode, or None if denied and raise_on_deny is False.

        Raises:
            PolicyViolationError: If insert is denied and raise_on_deny is True.
        """
        effective_ns = namespace or self._memory.namespace

        # Extract relations from graph
        relations = []
        if graph is not None:
            for _, _, data in graph.edges(data=True):
                rel = data.get("relation") or data.get("label") or data.get("type")
                if rel:
                    relations.append(str(rel))

        # Build evaluation context
        ctx = EvaluationContext(
            action=PolicyAction.INSERT,
            namespace=effective_ns,
            content=meta.get("text") or meta.get("content") if meta else None,
            metadata=meta or {},
            relations=relations,
            total_episodes=len(self._memory.episodes),
        )

        # Evaluate policy
        decision = self._engine.evaluate(ctx)

        # Handle decision
        if not self._handle_decision(decision):
            return None

        # Proceed with insert
        return self._memory.insert(vector, graph, meta, namespace=namespace)

    def summarise(self) -> None:
        """
        Trigger summarization with policy enforcement.

        Policies can control:
        - Whether summarization should happen
        - Depth of hierarchy
        - Fidelity requirements
        """
        ctx = EvaluationContext(
            action=PolicyAction.SUMMARIZE,
            namespace=self._memory.namespace,
            total_episodes=len(self._memory.episodes),
        )

        decision = self._engine.evaluate(ctx)

        if not self._handle_decision(decision):
            return

        # Apply policy modifications
        mods = decision.modifications
        if "depth" in mods:
            # Could be used to adjust summarization behavior
            pass

        self._memory.summarise()

    def apply_decay(self, cycles: int = 1) -> List[PolicyDecision]:
        """
        Apply trust decay to all episodes based on policies.

        Args:
            cycles: Number of decay cycles to apply.

        Returns:
            List of policy decisions (one per episode).
        """
        decisions = []

        for episode in list(self._memory.episodes):
            ctx = EvaluationContext(
                action=PolicyAction.DECAY,
                namespace=self._memory.namespace,
                episode_id=episode.id,
                current_trust=episode.trust,
                cycles_unused=cycles,
                metadata=episode.meta,
            )

            decision = self._engine.evaluate(ctx)
            decisions.append(decision)

            # Apply decay result
            mods = decision.modifications
            if mods.get("decayed", False):
                new_trust = mods.get("new_trust", episode.trust)
                old_trust = episode.trust
                episode.trust = new_trust

                # Emit decay event
                if self._emitter:
                    self._emitter.emit_decay(
                        namespace=self._memory.namespace,
                        episode_id=episode.id,
                        old_trust=old_trust,
                        new_trust=new_trust,
                        policy_name=decision.policy_name,
                    )

                # Check if should prune
                if mods.get("should_prune", False):
                    self._prune_episode(episode, decision)

        return decisions

    def _prune_episode(self, episode: Episode, decision: PolicyDecision) -> None:
        """Prune a single episode after policy approval."""
        # Emit prune event
        if self._emitter:
            self._emitter.emit_prune(
                namespace=self._memory.namespace,
                episode_id=episode.id,
                reason=f"Trust below threshold: {episode.trust:.3f}",
                policy_name=decision.policy_name,
            )

        # Remove from memory
        if episode in self._memory.episodes:
            self._memory.episodes.remove(episode)

        # Remove from store
        if self._memory.store is not None:
            try:
                self._memory.store.remove(episode.id)
            except Exception as e:
                logger.warning(f"Failed to remove episode from store: {e}")

        # Remove from indexes
        try:
            self._memory._index_remove_episode(episode.id)
        except Exception:
            pass

        try:
            self._memory.entity_index.remove(episode.id)
        except Exception:
            pass

    def prune_below_threshold(self, namespace: Optional[str] = None) -> List[str]:
        """
        Prune episodes below trust threshold with policy enforcement.

        Unlike prune_below_trust which uses a fixed threshold, this method
        uses the threshold from the applicable policy.

        Args:
            namespace: Optional namespace to prune.

        Returns:
            List of pruned episode IDs.
        """
        pruned = []
        effective_ns = namespace or self._memory.namespace

        for episode in list(self._memory.episodes):
            ctx = EvaluationContext(
                action=PolicyAction.PRUNE,
                namespace=effective_ns,
                episode_id=episode.id,
                current_trust=episode.trust,
                metadata=episode.meta,
                relations=[],  # Could extract from graph
                edge_count=episode.G.number_of_edges() if episode.G else 0,
            )

            decision = self._engine.evaluate(ctx)

            if decision.is_allowed():
                self._prune_episode(episode, decision)
                pruned.append(episode.id)

        return pruned

    def check_capacity(self) -> Optional[PolicyDecision]:
        """
        Check if memory is approaching capacity limits.

        Returns:
            PolicyDecision with warnings if approaching limit, None otherwise.
        """
        ctx = EvaluationContext(
            action=PolicyAction.INSERT,
            namespace=self._memory.namespace,
            total_episodes=len(self._memory.episodes),
        )

        decision = self._engine.evaluate(ctx)

        if decision.warnings:
            return decision

        return None

    # Delegate all other attributes to underlying memory
    def __getattr__(self, name: str) -> Any:
        """Delegate attribute access to underlying memory."""
        return getattr(self._memory, name)


def create_policy_memory(
    memory_kwargs: Optional[Dict[str, Any]] = None,
    policies: Optional[List[Any]] = None,  # List[MemoryPolicy]
    enable_logging: bool = True,
    log_file: Optional[str] = None,
) -> PolicyAwareMemory:
    """
    Factory function to create a policy-aware memory instance.

    Args:
        memory_kwargs: Kwargs to pass to Memory constructor.
        policies: List of MemoryPolicy objects to register.
        enable_logging: Whether to enable policy event logging.
        log_file: Optional file path for JSON Lines event log.

    Returns:
        Configured PolicyAwareMemory instance.
    """
    from hieromem_core.memory import Memory
    from hieromem_core.policy.events import (
        LoggingEventHandler,
        FileEventHandler,
        PolicyEventEmitter,
    )
    from hieromem_core.policy.schema import MemoryPolicy

    # Create memory
    memory = Memory(**(memory_kwargs or {}))

    # Create engine
    engine = PolicyEngine()

    # Register policies (or default)
    if policies:
        for policy in policies:
            engine.register_policy(policy)
    else:
        engine.register_policy(MemoryPolicy.default())

    # Create emitter
    emitter = None
    if enable_logging:
        emitter = PolicyEventEmitter()
        emitter.add_handler(LoggingEventHandler())

        if log_file:
            emitter.add_handler(FileEventHandler(log_file))

    return PolicyAwareMemory(
        memory=memory,
        policy_engine=engine,
        event_emitter=emitter,
    )
