"""
Agent Guardrails / Policy Constraints — User-defined behavioral boundaries.

No other framework lets users set behavioral constraints on learned policies.
Agent Lightning optimizes models with RL but has zero runtime constraints.

Users define rules like:
- "Never call this API more than 3x per task"
- "Always validate output format before returning"
- "Max 5 steps per task"
- "Never use tool X for task type Y"

These become part of policy recommendations and are enforced at runtime.
"""

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


class GuardrailType(str, Enum):
    """Types of guardrails users can define."""

    MAX_STEPS = "max_steps"  # Limit total steps per task
    MAX_TOOL_CALLS = "max_tool_calls"  # Limit calls to a specific tool
    MAX_TOKENS = "max_tokens"  # Token budget per task
    MAX_LATENCY_MS = "max_latency_ms"  # Time budget per task
    REQUIRED_STEP = "required_step"  # Must include this step/tool
    FORBIDDEN_TOOL = "forbidden_tool"  # Never use this tool for task_type
    OUTPUT_FORMAT = "output_format"  # Output must match pattern
    CUSTOM = "custom"  # Custom validation function


class GuardrailAction(str, Enum):
    """What to do when a guardrail is violated."""

    WARN = "warn"  # Log warning but continue
    BLOCK = "block"  # Stop execution immediately
    FALLBACK = "fallback"  # Use fallback behavior
    RETRY = "retry"  # Retry with guardrail-aware prompt


@dataclass
class Guardrail:
    """A single user-defined guardrail."""

    guardrail_id: str
    name: str
    guardrail_type: GuardrailType
    action: GuardrailAction = GuardrailAction.WARN

    # Scope
    task_type: Optional[str] = None  # None = applies to ALL tasks
    tool_name: Optional[str] = None  # For tool-specific guardrails

    # Parameters (type-dependent)
    max_value: Optional[int] = None  # For max_steps, max_tool_calls, etc.
    required_tool: Optional[str] = None  # For required_step
    output_pattern: Optional[str] = None  # Regex for output_format
    custom_message: Optional[str] = None  # Human-readable description

    # Runtime state
    is_active: bool = True
    violation_count: int = 0
    last_violation_at: Optional[datetime] = None
    created_at: datetime = field(default_factory=datetime.utcnow)

    def to_dict(self) -> Dict:
        return {
            "guardrail_id": self.guardrail_id,
            "name": self.name,
            "guardrail_type": self.guardrail_type.value,
            "action": self.action.value,
            "task_type": self.task_type,
            "tool_name": self.tool_name,
            "max_value": self.max_value,
            "required_tool": self.required_tool,
            "output_pattern": self.output_pattern,
            "custom_message": self.custom_message,
            "is_active": self.is_active,
            "violation_count": self.violation_count,
            "created_at": self.created_at.isoformat(),
        }


@dataclass
class GuardrailViolation:
    """Record of a guardrail being violated."""

    violation_id: str
    guardrail_id: str
    guardrail_name: str
    trajectory_id: str
    step_index: int
    violation_detail: str
    action_taken: str  # What action was taken
    timestamp: datetime = field(default_factory=datetime.utcnow)

    def to_dict(self) -> Dict:
        return {
            "violation_id": self.violation_id,
            "guardrail_id": self.guardrail_id,
            "guardrail_name": self.guardrail_name,
            "trajectory_id": self.trajectory_id,
            "step_index": self.step_index,
            "violation_detail": self.violation_detail,
            "action_taken": self.action_taken,
            "timestamp": self.timestamp.isoformat(),
        }


# =============================================================================
# Guardrail Engine
# =============================================================================


class GuardrailEngine:
    """
    Runtime guardrail enforcement engine.

    Checks guardrails at each step and before returning output.
    Can warn, block, fallback, or retry based on guardrail configuration.
    """

    def __init__(self):
        self._guardrails: Dict[str, Guardrail] = {}  # id → guardrail
        self._violations: List[GuardrailViolation] = []
        self._violation_counter = 0

        # Runtime tracking per trajectory
        self._trajectory_state: Dict[str, Dict] = {}
        # trajectory_id → {step_count, tool_counts, token_count, latency_ms, tools_used}

    # =========================================================================
    # Guardrail CRUD
    # =========================================================================

    def add_guardrail(self, guardrail: Guardrail) -> None:
        """Add a guardrail."""
        self._guardrails[guardrail.guardrail_id] = guardrail

    def remove_guardrail(self, guardrail_id: str) -> bool:
        """Remove a guardrail."""
        return self._guardrails.pop(guardrail_id, None) is not None

    def get_guardrails(self, task_type: str = None) -> List[Guardrail]:
        """Get active guardrails, optionally for a specific task type."""
        guardrails = [g for g in self._guardrails.values() if g.is_active]
        if task_type:
            guardrails = [g for g in guardrails if g.task_type is None or g.task_type == task_type]
        return guardrails

    def toggle_guardrail(self, guardrail_id: str, active: bool) -> bool:
        """Enable/disable a guardrail."""
        g = self._guardrails.get(guardrail_id)
        if g:
            g.is_active = active
            return True
        return False

    # =========================================================================
    # Runtime Tracking
    # =========================================================================

    def start_tracking(self, trajectory_id: str) -> None:
        """Start tracking for a new trajectory."""
        self._trajectory_state[trajectory_id] = {
            "step_count": 0,
            "tool_counts": {},  # tool_name → count
            "token_count": 0,
            "latency_ms": 0,
            "tools_used": set(),
        }

    def record_step(
        self,
        trajectory_id: str,
        step_index: int,
        tool_name: Optional[str],
        tokens_used: int = 0,
        latency_ms: int = 0,
    ) -> None:
        """Record a step for guardrail tracking."""
        state = self._trajectory_state.get(trajectory_id)
        if not state:
            self.start_tracking(trajectory_id)
            state = self._trajectory_state[trajectory_id]

        state["step_count"] += 1
        state["token_count"] += tokens_used
        state["latency_ms"] += latency_ms

        if tool_name:
            state["tool_counts"][tool_name] = state["tool_counts"].get(tool_name, 0) + 1
            state["tools_used"].add(tool_name)

    def end_tracking(self, trajectory_id: str) -> None:
        """Clean up tracking for a completed trajectory."""
        self._trajectory_state.pop(trajectory_id, None)

    # =========================================================================
    # Guardrail Checking
    # =========================================================================

    def check_pre_step(
        self,
        trajectory_id: str,
        step_index: int,
        task_type: str,
        tool_name: Optional[str] = None,
    ) -> List[GuardrailViolation]:
        """
        Check guardrails BEFORE a step executes.

        Returns list of violations (empty = all clear).
        """
        violations = []
        state = self._trajectory_state.get(trajectory_id, {})
        guardrails = self.get_guardrails(task_type)

        for g in guardrails:
            violation = self._check_guardrail_pre_step(g, state, step_index, tool_name)
            if violation:
                v = self._create_violation(g, trajectory_id, step_index, violation)
                violations.append(v)

        return violations

    def check_post_output(
        self,
        trajectory_id: str,
        task_type: str,
        output: Any,
    ) -> List[GuardrailViolation]:
        """
        Check guardrails AFTER output is produced.

        Checks output format, required steps, etc.
        """
        violations = []
        state = self._trajectory_state.get(trajectory_id, {})
        guardrails = self.get_guardrails(task_type)

        for g in guardrails:
            violation = self._check_guardrail_post_output(g, state, output)
            if violation:
                step_index = state.get("step_count", 0)
                v = self._create_violation(g, trajectory_id, step_index, violation)
                violations.append(v)

        return violations

    def get_prompt_constraints(self, task_type: str) -> str:
        """
        Generate constraint text to inject into agent prompt.

        This is how guardrails become part of policy recommendations.
        """
        guardrails = self.get_guardrails(task_type)
        if not guardrails:
            return ""

        constraints = []
        for g in guardrails:
            text = self._guardrail_to_prompt(g)
            if text:
                constraints.append(f"- {text}")

        if constraints:
            return "CONSTRAINTS (must follow):\n" + "\n".join(constraints)
        return ""

    # =========================================================================
    # Internal Checking Logic
    # =========================================================================

    def _check_guardrail_pre_step(
        self, g: Guardrail, state: Dict, step_index: int, tool_name: Optional[str]
    ) -> Optional[str]:
        """Check a single guardrail before step execution."""

        if g.guardrail_type == GuardrailType.MAX_STEPS:
            if g.max_value and state.get("step_count", 0) >= g.max_value:
                return f"Max steps ({g.max_value}) exceeded"

        elif g.guardrail_type == GuardrailType.MAX_TOOL_CALLS:
            if g.tool_name and tool_name == g.tool_name:
                current = state.get("tool_counts", {}).get(tool_name, 0)
                if g.max_value and current >= g.max_value:
                    return f"Max calls to '{tool_name}' ({g.max_value}) exceeded"

        elif g.guardrail_type == GuardrailType.MAX_TOKENS:
            if g.max_value and state.get("token_count", 0) >= g.max_value:
                return f"Token budget ({g.max_value}) exceeded"

        elif g.guardrail_type == GuardrailType.MAX_LATENCY_MS:
            if g.max_value and state.get("latency_ms", 0) >= g.max_value:
                return f"Latency budget ({g.max_value}ms) exceeded"

        elif g.guardrail_type == GuardrailType.FORBIDDEN_TOOL:
            if g.tool_name and tool_name == g.tool_name:
                return f"Forbidden tool '{tool_name}' used"

        return None

    def _check_guardrail_post_output(self, g: Guardrail, state: Dict, output: Any) -> Optional[str]:
        """Check a single guardrail after output is produced."""

        if g.guardrail_type == GuardrailType.REQUIRED_STEP:
            if g.required_tool and g.required_tool not in state.get("tools_used", set()):
                return f"Required tool '{g.required_tool}' was not used"

        elif g.guardrail_type == GuardrailType.OUTPUT_FORMAT:
            if g.output_pattern and output is not None:
                output_str = str(output)
                try:
                    if not re.search(g.output_pattern, output_str):
                        return f"Output does not match required format: {g.output_pattern}"
                except re.error:
                    pass  # Invalid regex, skip

        return None

    def _guardrail_to_prompt(self, g: Guardrail) -> Optional[str]:
        """Convert a guardrail to a prompt constraint."""
        if g.custom_message:
            return g.custom_message

        if g.guardrail_type == GuardrailType.MAX_STEPS:
            return f"Complete this task in {g.max_value} steps or fewer"
        elif g.guardrail_type == GuardrailType.MAX_TOOL_CALLS:
            return f"Do not call '{g.tool_name}' more than {g.max_value} times"
        elif g.guardrail_type == GuardrailType.MAX_TOKENS:
            return f"Stay within a {g.max_value} token budget"
        elif g.guardrail_type == GuardrailType.FORBIDDEN_TOOL:
            return f"Do not use the '{g.tool_name}' tool"
        elif g.guardrail_type == GuardrailType.REQUIRED_STEP:
            return f"You must use '{g.required_tool}' at some point during execution"
        elif g.guardrail_type == GuardrailType.OUTPUT_FORMAT:
            return f"Output must match format: {g.output_pattern}"

        return None

    def _create_violation(
        self, g: Guardrail, trajectory_id: str, step_index: int, detail: str
    ) -> GuardrailViolation:
        """Create and record a violation."""
        self._violation_counter += 1
        v = GuardrailViolation(
            violation_id=f"gv_{self._violation_counter}",
            guardrail_id=g.guardrail_id,
            guardrail_name=g.name,
            trajectory_id=trajectory_id,
            step_index=step_index,
            violation_detail=detail,
            action_taken=g.action.value,
        )
        g.violation_count += 1
        g.last_violation_at = datetime.utcnow()
        self._violations.append(v)
        return v

    # =========================================================================
    # Stats & Export
    # =========================================================================

    def get_stats(self) -> Dict:
        """Get guardrail stats for dashboard."""
        active = [g for g in self._guardrails.values() if g.is_active]
        return {
            "total_guardrails": len(self._guardrails),
            "active_guardrails": len(active),
            "total_violations": len(self._violations),
            "violations_by_type": self._violations_by_type(),
            "most_violated": self._most_violated(),
        }

    def _violations_by_type(self) -> Dict[str, int]:
        """Count violations by guardrail type."""
        counts: Dict[str, int] = {}
        for v in self._violations:
            g = self._guardrails.get(v.guardrail_id)
            if g:
                key = g.guardrail_type.value
                counts[key] = counts.get(key, 0) + 1
        return counts

    def _most_violated(self, limit: int = 5) -> List[Dict]:
        """Get most frequently violated guardrails."""
        sorted_g = sorted(
            self._guardrails.values(),
            key=lambda g: g.violation_count,
            reverse=True,
        )
        return [
            {"guardrail_id": g.guardrail_id, "name": g.name, "violations": g.violation_count}
            for g in sorted_g[:limit]
            if g.violation_count > 0
        ]

    def get_violations(self, trajectory_id: str = None, limit: int = 50) -> List[Dict]:
        """Get violations, optionally filtered by trajectory."""
        violations = self._violations
        if trajectory_id:
            violations = [v for v in violations if v.trajectory_id == trajectory_id]
        return [v.to_dict() for v in violations[-limit:]]

    def export_guardrails(self) -> List[Dict]:
        """Export guardrails for persistence."""
        return [g.to_dict() for g in self._guardrails.values()]

    def load_guardrails(self, guardrails_data: List[Dict]) -> int:
        """Load guardrails from server."""
        loaded = 0
        for data in guardrails_data:
            g = Guardrail(
                guardrail_id=data.get("guardrail_id", ""),
                name=data.get("name", ""),
                guardrail_type=GuardrailType(data.get("guardrail_type", "custom")),
                action=GuardrailAction(data.get("action", "warn")),
                task_type=data.get("task_type"),
                tool_name=data.get("tool_name"),
                max_value=data.get("max_value"),
                required_tool=data.get("required_tool"),
                output_pattern=data.get("output_pattern"),
                custom_message=data.get("custom_message"),
                is_active=data.get("is_active", True),
            )
            self._guardrails[g.guardrail_id] = g
            loaded += 1
        return loaded
