"""
Policy Schema
=============

Data models for HIEROMEM memory policies.

Policies control:
- Retention: How long memories live, decay curves
- Summarization: When and how to compress memories
- Graph: Edge TTL, orphan pruning, traversal limits
- Governance: Forbidden relations, capacity limits, trust bounds
- Priority: High/low priority patterns for selective retention
"""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import timedelta
from enum import Enum
from typing import Any, Dict, List, Optional, Set
import re


class DecayCurve(Enum):
    """Trust decay curve types."""
    LINEAR = "linear"  # trust -= decay_rate
    EXPONENTIAL = "exponential"  # trust *= decay_rate
    STEP = "step"  # trust = 0 after threshold
    NONE = "none"  # No decay


@dataclass
class RetentionPolicy:
    """Controls how long memories are retained and how they decay."""

    # Maximum age before forced removal (None = no limit)
    horizon: Optional[timedelta] = None

    # Decay curve type
    decay_curve: DecayCurve = DecayCurve.EXPONENTIAL

    # Base decay rate (lambda in the math)
    # For exponential: trust *= base_lambda each cycle
    # For linear: trust -= (1 - base_lambda) each cycle
    base_lambda: float = 0.95

    # Minimum trust threshold before pruning
    min_trust: float = 0.1

    # How often to run decay (in seconds, None = on access only)
    decay_interval: Optional[int] = None

    # Boost factor when memory is accessed
    access_boost: float = 1.0  # 1.0 = reset to full trust

    def __post_init__(self) -> None:
        if isinstance(self.decay_curve, str):
            self.decay_curve = DecayCurve(self.decay_curve)
        if isinstance(self.horizon, (int, float)):
            self.horizon = timedelta(days=self.horizon)

    def compute_decay(self, current_trust: float, cycles_unused: int = 1) -> float:
        """Compute new trust after decay cycles."""
        if self.decay_curve == DecayCurve.NONE:
            return current_trust

        trust = current_trust
        for _ in range(cycles_unused):
            if self.decay_curve == DecayCurve.EXPONENTIAL:
                trust *= self.base_lambda
            elif self.decay_curve == DecayCurve.LINEAR:
                trust -= (1.0 - self.base_lambda)
            elif self.decay_curve == DecayCurve.STEP:
                if trust < self.min_trust * 2:  # Near threshold
                    trust = 0.0
                    break

        return max(0.0, trust)


@dataclass
class SummarizationPolicy:
    """Controls when and how memories are summarized/compressed."""

    # Trigger summarization when episode count exceeds this
    max_episodes: Optional[int] = 100

    # Trigger summarization when oldest episode exceeds this age
    max_age: Optional[timedelta] = None

    # Number of hierarchy levels to maintain
    depth: int = 2

    # Preserve graph edges in summaries
    preserve_graph: bool = True

    # Minimum fidelity score for summaries (0-1)
    min_fidelity: float = 0.7

    # Batch size for summarization
    batch_size: int = 5

    def __post_init__(self) -> None:
        if isinstance(self.max_age, (int, float)):
            self.max_age = timedelta(days=self.max_age)


@dataclass
class GraphPolicy:
    """Controls graph edge behavior and traversal limits."""

    # Edge time-to-live (None = no expiration)
    edge_ttl: Optional[timedelta] = None

    # Remove nodes with no edges
    prune_orphans: bool = True

    # Maximum hops for graph traversal queries
    max_hops: int = 5

    # Maximum edges per node
    max_edges_per_node: Optional[int] = None

    # Edge types to always preserve (never prune)
    protected_edge_types: Set[str] = field(default_factory=set)

    def __post_init__(self) -> None:
        if isinstance(self.edge_ttl, (int, float)):
            self.edge_ttl = timedelta(days=self.edge_ttl)
        if isinstance(self.protected_edge_types, list):
            self.protected_edge_types = set(self.protected_edge_types)


@dataclass
class GovernancePolicy:
    """Safety and governance constraints."""

    # Relations that cannot be stored
    forbidden_relations: Set[str] = field(default_factory=set)

    # Maximum episodes per namespace (hard cap)
    max_episodes: Optional[int] = None

    # Maximum trust boost on access (prevents gaming)
    max_trust_boost: float = 1.5

    # Minimum lambda value (prevents too-fast decay)
    min_lambda: float = 0.8

    # Maximum lambda value (prevents never-decay)
    max_lambda: float = 0.99

    # Require audit logging for all operations
    audit_required: bool = False

    # Content patterns to reject (regex)
    blocked_content_patterns: List[str] = field(default_factory=list)

    # Compiled patterns (internal)
    _compiled_patterns: List[re.Pattern] = field(default_factory=list, repr=False)

    def __post_init__(self) -> None:
        if isinstance(self.forbidden_relations, list):
            self.forbidden_relations = set(self.forbidden_relations)
        # Compile regex patterns
        self._compiled_patterns = [
            re.compile(p, re.IGNORECASE)
            for p in self.blocked_content_patterns
        ]

    def is_relation_forbidden(self, relation: str) -> bool:
        """Check if a relation type is forbidden."""
        return relation.lower() in {r.lower() for r in self.forbidden_relations}

    def is_content_blocked(self, content: str) -> bool:
        """Check if content matches any blocked patterns."""
        return any(p.search(content) for p in self._compiled_patterns)


@dataclass
class PriorityPolicy:
    """Controls priority-based retention."""

    # Patterns that mark high-priority memories (slower decay)
    high_priority_patterns: List[str] = field(default_factory=list)

    # Patterns that mark low-priority memories (faster decay)
    low_priority_patterns: List[str] = field(default_factory=list)

    # Lambda multiplier for high-priority (> 1 = slower decay)
    high_priority_lambda_multiplier: float = 1.05

    # Lambda multiplier for low-priority (< 1 = faster decay)
    low_priority_lambda_multiplier: float = 0.9

    # Tags that mark memories as critical (never decay)
    critical_tags: Set[str] = field(default_factory=set)

    def __post_init__(self) -> None:
        if isinstance(self.critical_tags, list):
            self.critical_tags = set(self.critical_tags)

    def get_priority_multiplier(self, meta: Dict[str, Any]) -> float:
        """Get lambda multiplier based on metadata."""
        # Check critical tags first
        tags = meta.get("tags", [])
        if isinstance(tags, str):
            tags = [tags]
        if any(t in self.critical_tags for t in tags):
            return float("inf")  # Never decay

        content = str(meta.get("text", "")) + str(meta.get("content", ""))

        # Check high priority patterns
        for pattern in self.high_priority_patterns:
            if pattern.lower() in content.lower():
                return self.high_priority_lambda_multiplier

        # Check low priority patterns
        for pattern in self.low_priority_patterns:
            if pattern.lower() in content.lower():
                return self.low_priority_lambda_multiplier

        return 1.0  # Default: no modification


@dataclass
class MemoryPolicy:
    """
    Complete memory policy specification.

    Combines retention, summarization, graph, governance, and priority
    policies into a single configuration unit.
    """

    # Policy identifier
    name: str

    # Target namespace (or "*" for all)
    namespace: str = "*"

    # Sub-policies
    retention: RetentionPolicy = field(default_factory=RetentionPolicy)
    summarization: SummarizationPolicy = field(default_factory=SummarizationPolicy)
    graph: GraphPolicy = field(default_factory=GraphPolicy)
    governance: GovernancePolicy = field(default_factory=GovernancePolicy)
    priority: PriorityPolicy = field(default_factory=PriorityPolicy)

    # Policy metadata
    version: str = "1.0"
    description: str = ""
    enabled: bool = True

    @classmethod
    def default(cls) -> "MemoryPolicy":
        """Create a default policy with sensible defaults."""
        return cls(
            name="default",
            namespace="*",
            retention=RetentionPolicy(
                horizon=timedelta(days=30),
                decay_curve=DecayCurve.EXPONENTIAL,
                base_lambda=0.95,
                min_trust=0.1,
            ),
            summarization=SummarizationPolicy(
                max_episodes=100,
                depth=2,
                preserve_graph=True,
            ),
            graph=GraphPolicy(
                edge_ttl=timedelta(days=60),
                prune_orphans=True,
                max_hops=5,
            ),
            governance=GovernancePolicy(
                forbidden_relations={"kills", "harms", "threatens"},
                max_episodes=10000,
            ),
            description="Default memory policy with balanced retention and governance",
        )

    @classmethod
    def aggressive_decay(cls, name: str = "aggressive") -> "MemoryPolicy":
        """Create a policy with aggressive decay for transient data."""
        return cls(
            name=name,
            retention=RetentionPolicy(
                horizon=timedelta(days=7),
                decay_curve=DecayCurve.EXPONENTIAL,
                base_lambda=0.85,
                min_trust=0.2,
            ),
            summarization=SummarizationPolicy(
                max_episodes=50,
                depth=1,
            ),
            description="Aggressive decay for transient/chat data",
        )

    @classmethod
    def long_term(cls, name: str = "long_term") -> "MemoryPolicy":
        """Create a policy for long-term retention."""
        return cls(
            name=name,
            retention=RetentionPolicy(
                horizon=timedelta(days=365),
                decay_curve=DecayCurve.EXPONENTIAL,
                base_lambda=0.99,
                min_trust=0.05,
            ),
            summarization=SummarizationPolicy(
                max_episodes=1000,
                depth=3,
                preserve_graph=True,
            ),
            description="Long-term retention for important knowledge",
        )

    @classmethod
    def safety_critical(cls, name: str = "safety_critical") -> "MemoryPolicy":
        """Create a policy for safety-critical data that never decays."""
        return cls(
            name=name,
            retention=RetentionPolicy(
                horizon=None,  # No expiration
                decay_curve=DecayCurve.NONE,
                base_lambda=1.0,
            ),
            governance=GovernancePolicy(
                audit_required=True,
                max_trust_boost=1.0,  # No boosting
            ),
            description="Safety-critical data with no decay and full audit",
        )

    def to_dict(self) -> Dict[str, Any]:
        """Convert policy to dictionary for serialization."""
        return {
            "apiVersion": "hieromem/v1",
            "kind": "MemoryPolicy",
            "metadata": {
                "name": self.name,
                "namespace": self.namespace,
                "version": self.version,
                "description": self.description,
            },
            "spec": {
                "enabled": self.enabled,
                "retention": {
                    "horizon": self.retention.horizon.days if self.retention.horizon else None,
                    "decay_curve": self.retention.decay_curve.value,
                    "base_lambda": self.retention.base_lambda,
                    "min_trust": self.retention.min_trust,
                    "decay_interval": self.retention.decay_interval,
                    "access_boost": self.retention.access_boost,
                },
                "summarization": {
                    "max_episodes": self.summarization.max_episodes,
                    "max_age": self.summarization.max_age.days if self.summarization.max_age else None,
                    "depth": self.summarization.depth,
                    "preserve_graph": self.summarization.preserve_graph,
                    "min_fidelity": self.summarization.min_fidelity,
                    "batch_size": self.summarization.batch_size,
                },
                "graph": {
                    "edge_ttl": self.graph.edge_ttl.days if self.graph.edge_ttl else None,
                    "prune_orphans": self.graph.prune_orphans,
                    "max_hops": self.graph.max_hops,
                    "max_edges_per_node": self.graph.max_edges_per_node,
                    "protected_edge_types": list(self.graph.protected_edge_types),
                },
                "governance": {
                    "forbidden_relations": list(self.governance.forbidden_relations),
                    "max_episodes": self.governance.max_episodes,
                    "max_trust_boost": self.governance.max_trust_boost,
                    "min_lambda": self.governance.min_lambda,
                    "max_lambda": self.governance.max_lambda,
                    "audit_required": self.governance.audit_required,
                    "blocked_content_patterns": self.governance.blocked_content_patterns,
                },
                "priority": {
                    "high_priority_patterns": self.priority.high_priority_patterns,
                    "low_priority_patterns": self.priority.low_priority_patterns,
                    "high_priority_lambda_multiplier": self.priority.high_priority_lambda_multiplier,
                    "low_priority_lambda_multiplier": self.priority.low_priority_lambda_multiplier,
                    "critical_tags": list(self.priority.critical_tags),
                },
            },
        }
