"""Autonomy guardrail enforcement service.

Provides a high-level service interface for checking autonomy
guardrails during plan execution and recording enforcement
events to the plan's audit trail.

The service integrates with the plan metadata to persist audit
trail entries, providing a complete record of guardrail decisions.

Thread-safe: all state mutations are protected by a reentrant lock
so the service can be shared across threads (e.g. concurrent plan
step evaluations).
"""

from __future__ import annotations

import logging
import threading
from typing import Any

from cleveragents.application.services.cost_budget_service import CostBudgetService
from cleveragents.domain.models.core.autonomy_guardrails import (
    AutonomyGuardrails,
    GuardrailAuditEntry,
    GuardrailAuditTrail,
    GuardrailEventType,
    GuardrailResult,
)
from cleveragents.domain.models.core.cost_budget import BudgetCheckResult, BudgetLevel

logger = logging.getLogger(__name__)

# Key used in plan metadata to store the audit trail
_AUDIT_TRAIL_KEY = "guardrail_audit_trail"

# Key used in plan metadata to store guardrail state
_GUARDRAILS_KEY = "autonomy_guardrails"

# Hard ceiling for entries accepted through load_from_metadata to
# prevent memory exhaustion from untrusted or corrupted payloads.
_MAX_METADATA_ENTRIES = 50_000

# Hard ceiling for required_confirmations list length.
_MAX_CONFIRMATIONS = 500


class AutonomyGuardrailService:
    """Service for enforcing autonomy guardrails during plan execution.

    Manages guardrail state and audit trail persistence via plan
    metadata dictionaries. Each plan can have its own guardrail
    configuration and audit trail stored in its metadata.

    All public methods are thread-safe.
    """

    def __init__(
        self,
        cost_budget_service: CostBudgetService | None = None,
    ) -> None:
        """Initialize the autonomy guardrail service.

        Args:
            cost_budget_service: Optional budget service for per-session
                and per-org budget hierarchy checks.
        """
        self._guardrails: dict[str, AutonomyGuardrails] = {}
        self._audit_trails: dict[str, GuardrailAuditTrail] = {}
        self._plan_session_map: dict[str, str] = {}
        self._cost_budget_service = cost_budget_service
        self._lock = threading.RLock()

    def configure_guardrails(
        self,
        plan_id: str,
        guardrails: AutonomyGuardrails,
    ) -> None:
        """Configure guardrails for a plan.

        Args:
            plan_id: The plan identifier.
            guardrails: The guardrail configuration to apply.
        """
        if not plan_id:
            raise ValueError("plan_id must not be empty")
        with self._lock:
            self._guardrails[plan_id] = guardrails
            if plan_id not in self._audit_trails:
                self._audit_trails[plan_id] = GuardrailAuditTrail()
        logger.debug("Configured guardrails for plan %s", plan_id)

    def get_guardrails(self, plan_id: str) -> AutonomyGuardrails | None:
        """Retrieve the guardrail configuration for a plan.

        Args:
            plan_id: The plan identifier.

        Returns:
            The guardrail configuration, or None if not configured.
        """
        with self._lock:
            return self._guardrails.get(plan_id)

    def check_step_limit(
        self,
        plan_id: str,
        current_step: int,
    ) -> bool:
        """Check if the plan has exceeded its maximum step limit.

        Args:
            plan_id: The plan identifier.
            current_step: The current step number.

        Returns:
            True if the step is allowed, False if blocked.
        """
        with self._lock:
            guardrails = self._guardrails.get(plan_id)
            if guardrails is None:
                return True

            guardrails.step_count = current_step
            allowed, reason = guardrails.check_step_limit()

            event_type = (
                GuardrailEventType.STEP_ALLOWED
                if allowed
                else GuardrailEventType.STEP_BLOCKED
            )
            result = GuardrailResult.ALLOWED if allowed else GuardrailResult.DENIED

            entry = GuardrailAuditEntry(
                event_type=event_type,
                guard_name="step_limit",
                result=result,
                reason=reason,
                context={
                    "current_step": current_step,
                    "max_steps": guardrails.max_steps,
                },
            )
            self._record_enforcement_locked(plan_id, entry)

            if not allowed:
                logger.warning(
                    "Step limit exceeded for plan %s: %s",
                    plan_id,
                    reason,
                )

            return allowed

    def check_tool_budget(
        self,
        plan_id: str,
        tool_cost: float,
    ) -> bool:
        """Check if the plan's budget allows a tool invocation.

        When allowed, the cost is recorded **before** the audit entry
        so that the trail reflects the post-deduction budget state.

        Args:
            plan_id: The plan identifier.
            tool_cost: The cost of the proposed tool invocation.

        Returns:
            True if the invocation is within budget, False if blocked.
        """
        with self._lock:
            guardrails = self._guardrails.get(plan_id)
            if guardrails is None:
                return True

            allowed, reason = guardrails.check_tool_budget(tool_cost)

            # Record cost *before* audit entry so context is accurate
            if allowed:
                guardrails.record_cost(tool_cost)

            event_type = (
                GuardrailEventType.BUDGET_ALLOWED
                if allowed
                else GuardrailEventType.BUDGET_BLOCKED
            )
            result = GuardrailResult.ALLOWED if allowed else GuardrailResult.DENIED

            entry = GuardrailAuditEntry(
                event_type=event_type,
                guard_name="tool_budget",
                result=result,
                reason=reason,
                context={
                    "tool_cost": tool_cost,
                    "budget_spent": guardrails.budget_spent,
                    "tool_budget": guardrails.tool_budget,
                },
            )
            self._record_enforcement_locked(plan_id, entry)

            if not allowed:
                logger.warning(
                    "Tool budget exceeded for plan %s: %s",
                    plan_id,
                    reason,
                )

            return allowed

    def check_wall_clock(self, plan_id: str) -> bool:
        """Check if the plan has exceeded its wall-clock time limit.

        Args:
            plan_id: The plan identifier.

        Returns:
            True if the plan is within the time limit, False if blocked.
        """
        with self._lock:
            guardrails = self._guardrails.get(plan_id)
            if guardrails is None:
                return True

            allowed, reason = guardrails.check_wall_clock()

            event_type = (
                GuardrailEventType.TIME_ALLOWED
                if allowed
                else GuardrailEventType.TIME_BLOCKED
            )
            result = GuardrailResult.ALLOWED if allowed else GuardrailResult.DENIED

            entry = GuardrailAuditEntry(
                event_type=event_type,
                guard_name="wall_clock",
                result=result,
                reason=reason,
                context={
                    "max_wall_clock_seconds": guardrails.max_wall_clock_seconds,
                },
            )
            self._record_enforcement_locked(plan_id, entry)

            if not allowed:
                logger.warning(
                    "Wall-clock limit exceeded for plan %s: %s",
                    plan_id,
                    reason,
                )

            return allowed

    def check_actor_tool_calls(
        self,
        plan_id: str,
        current_calls: int,
    ) -> bool:
        """Check if the actor has exceeded its per-invocation tool call limit.

        Args:
            plan_id: The plan identifier.
            current_calls: Number of tool calls made so far.

        Returns:
            True if allowed, False if blocked.
        """
        with self._lock:
            guardrails = self._guardrails.get(plan_id)
            if guardrails is None:
                return True

            allowed, reason = guardrails.check_actor_tool_calls(current_calls)

            event_type = (
                GuardrailEventType.ACTOR_LIMIT_ALLOWED
                if allowed
                else GuardrailEventType.ACTOR_LIMIT_BLOCKED
            )
            result = GuardrailResult.ALLOWED if allowed else GuardrailResult.DENIED

            entry = GuardrailAuditEntry(
                event_type=event_type,
                guard_name="actor_tool_calls",
                result=result,
                reason=reason,
                context={
                    "current_calls": current_calls,
                    "max_tool_calls_per_invocation": (
                        guardrails.actor_limits.max_tool_calls_per_invocation
                    ),
                },
            )
            self._record_enforcement_locked(plan_id, entry)

            if not allowed:
                logger.warning(
                    "Actor tool-call limit exceeded for plan %s: %s",
                    plan_id,
                    reason,
                )

            return allowed

    def check_retries_per_failure(
        self,
        plan_id: str,
        current_retries: int,
    ) -> bool:
        """Check if the retry count exceeds the per-failure retry limit.

        Args:
            plan_id: The plan identifier.
            current_retries: Number of retries attempted so far.

        Returns:
            True if the retry is allowed, False if blocked.
        """
        with self._lock:
            guardrails = self._guardrails.get(plan_id)
            if guardrails is None:
                return True

            allowed, reason = guardrails.check_retries_per_failure(current_retries)

            event_type = (
                GuardrailEventType.RETRY_ALLOWED
                if allowed
                else GuardrailEventType.RETRY_BLOCKED
            )
            result = GuardrailResult.ALLOWED if allowed else GuardrailResult.DENIED

            entry = GuardrailAuditEntry(
                event_type=event_type,
                guard_name="retries_per_failure",
                result=result,
                reason=reason,
                context={
                    "current_retries": current_retries,
                    "max_retries_per_failure": (
                        guardrails.actor_limits.max_retries_per_failure
                    ),
                },
            )
            self._record_enforcement_locked(plan_id, entry)

            if not allowed:
                logger.warning(
                    "Retry limit exceeded for plan %s: %s",
                    plan_id,
                    reason,
                )

            return allowed

    def check_confirmation_required(
        self,
        plan_id: str,
        operation: str,
    ) -> bool:
        """Check if an operation requires human confirmation.

        Args:
            plan_id: The plan identifier.
            operation: The operation name to check.

        Returns:
            True if confirmation is required, False otherwise.
        """
        with self._lock:
            guardrails = self._guardrails.get(plan_id)
            if guardrails is None:
                return False

            needs_confirmation, reason = guardrails.check_confirmation_required(
                operation,
            )

            if needs_confirmation:
                event_type = GuardrailEventType.CONFIRMATION_REQUIRED
                result = GuardrailResult.DENIED
            else:
                event_type = GuardrailEventType.CONFIRMATION_GRANTED
                result = GuardrailResult.ALLOWED

            entry = GuardrailAuditEntry(
                event_type=event_type,
                guard_name="confirmation",
                result=result,
                reason=reason,
                context={"operation": operation},
            )
            self._record_enforcement_locked(plan_id, entry)

            if needs_confirmation:
                logger.info(
                    "Confirmation required for plan %s operation '%s': %s",
                    plan_id,
                    operation,
                    reason,
                )

            return needs_confirmation

    def record_enforcement(
        self,
        plan_id: str,
        entry: GuardrailAuditEntry,
    ) -> None:
        """Persist an audit entry to the plan's audit trail.

        Args:
            plan_id: The plan identifier.
            entry: The audit entry to record.
        """
        with self._lock:
            self._record_enforcement_locked(plan_id, entry)

    def get_audit_trail(self, plan_id: str) -> GuardrailAuditTrail:
        """Retrieve the audit trail for a plan.

        Args:
            plan_id: The plan identifier.

        Returns:
            The audit trail for the plan. Returns an empty trail
            if no events have been recorded.
        """
        with self._lock:
            return self._audit_trails.get(plan_id, GuardrailAuditTrail())

    def get_audit_trail_as_metadata(
        self,
        plan_id: str,
    ) -> dict[str, Any]:
        """Serialize the audit trail for persistence in plan metadata.

        Args:
            plan_id: The plan identifier.

        Returns:
            A dictionary suitable for storing in plan metadata.
        """
        with self._lock:
            trail = self._audit_trails.get(plan_id, GuardrailAuditTrail())
            guardrails = self._guardrails.get(plan_id)

            metadata: dict[str, Any] = {
                _AUDIT_TRAIL_KEY: trail.model_dump(),
            }
            if guardrails is not None:
                metadata[_GUARDRAILS_KEY] = guardrails.model_dump()
            return metadata

    def load_from_metadata(
        self,
        plan_id: str,
        metadata: dict[str, Any],
    ) -> None:
        """Restore guardrail state from plan metadata.

        Applies size guards to prevent memory exhaustion from
        oversized or malicious payloads.

        Args:
            plan_id: The plan identifier.
            metadata: The plan metadata dictionary.

        Raises:
            ValueError: If the metadata contains oversized payloads.
        """
        with self._lock:
            if _GUARDRAILS_KEY in metadata:
                raw_guardrails = metadata[_GUARDRAILS_KEY]
                if isinstance(raw_guardrails, dict):
                    confirmations = raw_guardrails.get(
                        "required_confirmations",
                        [],
                    )
                    if (
                        isinstance(confirmations, list)
                        and len(confirmations) > _MAX_CONFIRMATIONS
                    ):
                        raise ValueError(
                            f"required_confirmations exceeds maximum of "
                            f"{_MAX_CONFIRMATIONS} entries"
                        )
                self._guardrails[plan_id] = AutonomyGuardrails.model_validate(
                    raw_guardrails,
                )

            if _AUDIT_TRAIL_KEY in metadata:
                raw_trail = metadata[_AUDIT_TRAIL_KEY]
                if isinstance(raw_trail, dict):
                    entries = raw_trail.get("entries", [])
                    if (
                        isinstance(entries, list)
                        and len(entries) > _MAX_METADATA_ENTRIES
                    ):
                        raise ValueError(
                            f"Audit trail exceeds maximum of "
                            f"{_MAX_METADATA_ENTRIES} entries"
                        )
                self._audit_trails[plan_id] = GuardrailAuditTrail.model_validate(
                    raw_trail
                )

    def associate_plan_with_session(
        self,
        plan_id: str,
        session_id: str,
    ) -> None:
        """Link a plan to a session for budget hierarchy checks.

        Args:
            plan_id: The plan identifier.
            session_id: The session identifier.

        Raises:
            ValueError: If *plan_id* or *session_id* is empty.
        """
        if not plan_id:
            raise ValueError("plan_id must not be empty")
        if not session_id:
            raise ValueError("session_id must not be empty")
        with self._lock:
            self._plan_session_map[plan_id] = session_id

    def check_budget_hierarchy(
        self,
        plan_id: str,
        tool_cost: float,
    ) -> BudgetCheckResult:
        """Check plan-level guardrails then delegate to CostBudgetService.

        Evaluation order:
        1. Plan-level tool budget (from AutonomyGuardrails)
        2. Session / Org budgets (from CostBudgetService)

        Args:
            plan_id: The plan identifier.
            tool_cost: The proposed tool invocation cost.

        Returns:
            :class:`BudgetCheckResult` with the enforcement outcome.

        Raises:
            ValueError: If *plan_id* is empty or *tool_cost* is negative.
            TypeError: If *tool_cost* is not numeric.
        """
        if not isinstance(tool_cost, (int, float)):
            raise TypeError("tool_cost must be a number")
        if not plan_id:
            raise ValueError("plan_id must not be empty")
        if tool_cost < 0:
            raise ValueError("tool_cost must be non-negative")

        with self._lock:
            # --- Plan-level check ---
            guardrails = self._guardrails.get(plan_id)
            if guardrails is not None:
                plan_budget = guardrails.tool_budget
                if (
                    plan_budget is not None
                    and guardrails.budget_spent + tool_cost > plan_budget
                ):
                    return BudgetCheckResult(
                        allowed=False,
                        exceeded_level=BudgetLevel.PLAN,
                        reason=(
                            f"Plan tool budget exceeded: "
                            f"{guardrails.budget_spent + tool_cost:.2f}"
                            f" > {plan_budget:.2f}"
                        ),
                    )

            # --- Session / Org hierarchy via CostBudgetService ---
            session_id = self._plan_session_map.get(plan_id)
            if session_id is not None and self._cost_budget_service is not None:
                return self._cost_budget_service.check_budget_hierarchy(
                    session_id=session_id,
                    plan_cost=tool_cost,
                )

            return BudgetCheckResult(allowed=True)

    def record_plan_cost_to_session(
        self,
        plan_id: str,
        cost: float,
    ) -> None:
        """Record cost against both the plan guardrails and the session.

        Args:
            plan_id: The plan identifier.
            cost: The cost to record.

        Raises:
            ValueError: If *plan_id* is empty or *cost* is negative.
            TypeError: If *cost* is not numeric.
        """
        if not isinstance(cost, (int, float)):
            raise TypeError("cost must be a number")
        if not plan_id:
            raise ValueError("plan_id must not be empty")
        if cost < 0:
            raise ValueError("cost must be non-negative")

        with self._lock:
            guardrails = self._guardrails.get(plan_id)
            if guardrails is not None:
                guardrails.record_cost(cost)

            session_id = self._plan_session_map.get(plan_id)
            if session_id is not None and self._cost_budget_service is not None:
                self._cost_budget_service.record_plan_cost(session_id, cost)

    def remove_plan(self, plan_id: str) -> None:
        """Remove all guardrail state for a plan.

        Args:
            plan_id: The plan identifier.
        """
        with self._lock:
            self._guardrails.pop(plan_id, None)
            self._audit_trails.pop(plan_id, None)
            self._plan_session_map.pop(plan_id, None)

    # ------------------------------------------------------------------
    # Private helpers
    # ------------------------------------------------------------------

    def _record_enforcement_locked(
        self,
        plan_id: str,
        entry: GuardrailAuditEntry,
    ) -> None:
        """Record an audit entry (caller must hold ``_lock``)."""
        if plan_id not in self._audit_trails:
            self._audit_trails[plan_id] = GuardrailAuditTrail()
        self._audit_trails[plan_id].add_entry(entry)
        logger.debug(
            "Recorded guardrail event for plan %s: %s (%s)",
            plan_id,
            entry.event_type.value,
            entry.result.value,
        )
