# Copyright (c) 2026 EaseTech Ltd.
# Licensed under the MIT License.
# https://sentinely.ai

"""sentinely — runtime guardrails for AI agents."""

from __future__ import annotations

import asyncio
import logging
import os
from datetime import datetime, timezone
from uuid import uuid4

logger = logging.getLogger("sentinely")

import httpx

from sentinely import scanner as _scanner_mod
from sentinely.circuit_breaker import (
    CircuitBreakerConfig,
    CircuitBreakerState,
    CircuitDecision,
    decide,
)
from sentinely.events import EventEmitter, SentinelyEvent
from sentinely.exceptions import SentinelyBlockedError, SentinelyConfigError, SentinelyQuarantineError
from sentinely.fingerprint import BehavioralFingerprint
from sentinely.intent import (
    ActionEvent,
    IntentScore,
    ScoringContext,
    compute_task_hash,
    score_intent,
)
from sentinely.memory import MemoryFirewall
from sentinely.multi_agent import AgentMessage, MultiAgentTracker


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

# Tool-name fragments that indicate an inter-agent call (used in score_action).
_INTER_AGENT_PATTERNS = frozenset(
    {"agent", "message", "invoke", "dispatch", "call_tool", "send_to"}
)


def protect(
    agent: object,
    task: str,
    policy: str = "strict",
    agent_id: str | None = None,
    org_id: str | None = None,
    tracker: MultiAgentTracker | None = None,
) -> ProtectedAgent:
    """Wrap *agent* with full Sentinely protection — the 3-line install.

    Parameters
    ----------
    tracker:
        Optional :class:`~sentinely.multi_agent.MultiAgentTracker` instance.
        When provided, any tool call whose name contains an inter-agent
        pattern (e.g. ``"call_agent"``, ``"send_message"``) is routed
        through ``tracker.track_message()`` in addition to the normal
        intent-scoring pipeline.
    """
    api_key = os.getenv("SENTINELY_API_KEY") or None

    agent_id_resolved = agent_id or f"agent_{uuid4().hex[:8]}"
    mode = "" if api_key else " (local mode)"
    print(
        f"\u2b21 Sentinely active{mode} \u2014 "
        f"agent: {agent_id_resolved}, policy: {policy}"
    )

    return ProtectedAgent(
        agent=agent,
        task=task,
        policy=policy,
        agent_id=agent_id_resolved,
        org_id=org_id or "local",
        tracker=tracker,
    )


# ---------------------------------------------------------------------------
# ProtectedAgent
# ---------------------------------------------------------------------------

class ProtectedAgent:
    """Wraps an agent executor with full Sentinely protection.

    Supports: ``.invoke()``, ``.ainvoke()``, ``.run()``, ``__call__()``.
    All of these delegate to the same internal ``_protected_invoke()``.
    """

    def __init__(
        self,
        agent: object,
        task: str,
        policy: str,
        agent_id: str,
        org_id: str,
        tracker: MultiAgentTracker | None = None,
    ) -> None:
        self._agent = agent
        self._task = task
        self._task_hash = compute_task_hash(task)
        self._policy = policy
        self._agent_id = agent_id
        self._org_id = org_id
        self._session_id = f"session_{uuid4().hex[:8]}"
        self._action_chain: list[dict] = []  # last 10 actions
        self._chain_lock = asyncio.Lock()

        api_key = os.getenv("SENTINELY_API_KEY")
        api_url = os.getenv("SENTINELY_API_URL", "https://api.sentinely.ai")
        env = os.getenv("SENTINELY_ENV", "development")

        self._api_key = api_key
        self._api_url = api_url
        self._emitter = EventEmitter(api_key, api_url, env)
        self._circuit_config = CircuitBreakerConfig(policy=policy)
        self._circuit_state = CircuitBreakerState()
        self._memory_firewall = MemoryFirewall()
        self._scanner = _scanner_mod

        # Behavioral fingerprinting (salami-slicing detection).
        self._fingerprint = BehavioralFingerprint(
            agent_id=agent_id,
            org_id=org_id,
        )
        self._salami_previously_detected: bool = False

        # Optional multi-agent tracker for inter-agent call monitoring.
        self._tracker: MultiAgentTracker | None = tracker

    # ------------------------------------------------------------------
    # Core invoke path
    # ------------------------------------------------------------------

    async def _protected_invoke(self, input_data: object, **kwargs: object) -> object:
        # Step 1: scan string input before the agent sees it.
        if isinstance(input_data, str):
            scan_result = await self._scanner.scan(input_data, use_llm=False)
            if not scan_result.is_safe and scan_result.risk_score >= 80:
                await self._emit_event(
                    "injection_detected",
                    None,
                    {},
                    scan_result.risk_score,
                    (
                        scan_result.threats_found[0]
                        if scan_result.threats_found
                        else "injection detected"
                    ),
                    "injection",
                    "block",
                )
                # Monitor policy: log but never block.
                if self._policy != "monitor":
                    raise SentinelyBlockedError(
                        risk_score=scan_result.risk_score,
                        reason=f"Input injection detected: {scan_result.threats_found}",
                        attack_type="injection",
                        action={"type": "input", "content": str(input_data)[:100]},
                    )
            input_data = scan_result.sanitized_content

        # Step 2: run the actual agent.
        return await self._invoke_agent(input_data, **kwargs)

    async def _invoke_agent(self, input_data: object, **kwargs: object) -> object:
        """Call whichever method the wrapped agent exposes."""
        if hasattr(self._agent, "ainvoke"):
            return await self._agent.ainvoke(input_data, **kwargs)  # type: ignore[union-attr]
        if hasattr(self._agent, "invoke"):
            return self._agent.invoke(input_data, **kwargs)  # type: ignore[union-attr]
        if hasattr(self._agent, "run"):
            return self._agent.run(input_data, **kwargs)  # type: ignore[union-attr]
        # Fall back to __call__.
        return await self._agent(input_data, **kwargs)  # type: ignore[operator]

    # ------------------------------------------------------------------
    # Tool-level scoring hook (called by adapters / frameworks)
    # ------------------------------------------------------------------

    async def score_action(
        self,
        tool_name: str,
        tool_params: dict,
        description: str = "",
    ) -> CircuitDecision:
        """Score a pending tool call and raise if blocked."""
        action = ActionEvent(
            tool_name=tool_name,
            tool_params=tool_params,
            description=description or tool_name,
        )
        async with self._chain_lock:
            chain_snapshot = [
                ActionEvent(**a) for a in self._action_chain[-10:]
            ]
        context = ScoringContext(
            original_task=self._task,
            task_hash=self._task_hash,
            agent_id=self._agent_id,
            action_chain=chain_snapshot,
            current_action=action,
            session_id=self._session_id,
        )

        intent_score = await score_intent(
            context, api_key=self._api_key, api_base_url=self._api_url,
        )
        decision = decide(intent_score, self._circuit_config)

        # Track consecutive flags.
        await self._circuit_state.record_decision(self._session_id, decision)
        if await self._circuit_state.should_auto_block(
            self._session_id, self._circuit_config
        ):
            decision = decision.model_copy(
                update={
                    "action": "block",
                    "reason": (
                        f"Auto-blocked: {self._circuit_config.max_consecutive_flags} "
                        "consecutive flags"
                    ),
                    "requires_human_approval": True,
                }
            )

        # Emit event (fire-and-forget).
        _ACTION_EVENT_TYPES = {
            "allow": "action_allowed",
            "flag": "action_flagged",
            "block": "action_blocked",
        }
        event_type = _ACTION_EVENT_TYPES.get(
            decision.action, f"action_{decision.action}"
        )
        asyncio.ensure_future(
            self._emit_event(
                event_type,
                tool_name,
                tool_params,
                intent_score.risk_score,
                intent_score.reason,
                intent_score.attack_type,
                decision.action,
            )
        )

        # ----------------------------------------------------------------
        # Update behavioral fingerprint — salami-slicing detection.
        # ----------------------------------------------------------------
        drift_report = self._fingerprint.update(
            risk_score=intent_score.risk_score,
            tool_name=tool_name,
            attack_type=intent_score.attack_type,
            session_id=self._session_id,
        )

        # Emit a one-time drift_detected event when salami is first confirmed.
        if drift_report.salami_detected and not self._salami_previously_detected:
            self._salami_previously_detected = True
            dv = drift_report.drift_vector
            asyncio.ensure_future(
                self._emit_event(
                    "drift_detected",
                    tool_name,
                    tool_params,
                    intent_score.risk_score,
                    (
                        f"Salami-slicing attack confirmed: risk elevated by "
                        f"+{dv.risk_elevation:.1f} points above baseline "
                        f"for {dv.consecutive_above_threshold} consecutive actions"
                    ),
                    "drift",
                    "flag",
                )
            )

        # Post fingerprint snapshot to API every 5 actions, or immediately
        # when salami is first detected.  Fire-and-forget — never blocks the
        # agent.
        if (
            self._fingerprint.total_actions % 5 == 0
            or drift_report.salami_detected
        ):
            asyncio.ensure_future(self._post_fingerprint())

        # If a multi-agent tracker is registered and this tool call looks like
        # an inter-agent communication, route it through the tracker so
        # behavioral drift can be detected across agent boundaries.
        if self._tracker is not None and any(
            pat in tool_name.lower() for pat in _INTER_AGENT_PATTERNS
        ):
            target_id = (
                tool_params.get("agent_id")
                or tool_params.get("to_agent_id")
                or tool_params.get("to")
                or "unknown"
            )
            content = str(
                tool_params.get("message")
                or tool_params.get("content")
                or description
            )
            inter_msg = AgentMessage(
                from_agent_id=self._agent_id,
                to_agent_id=str(target_id),
                content=content,
                message_type="agent",
            )
            asyncio.ensure_future(self._tracker.track_message(inter_msg))

        if decision.action == "block":
            raise SentinelyBlockedError(
                risk_score=intent_score.risk_score,
                reason=intent_score.reason,
                attack_type=intent_score.attack_type,
                action={"tool": tool_name, "params": tool_params},
                decision=decision,
            )

        # Record in action chain (only non-blocked actions).
        async with self._chain_lock:
            self._action_chain.append(action.model_dump())
            if len(self._action_chain) > 10:
                self._action_chain.pop(0)

        return decision

    # ------------------------------------------------------------------
    # Proxy methods — all delegate to _protected_invoke
    # ------------------------------------------------------------------

    async def invoke(self, input_data: object, **kwargs: object) -> object:
        return await self._protected_invoke(input_data, **kwargs)

    async def ainvoke(self, input_data: object, **kwargs: object) -> object:
        return await self._protected_invoke(input_data, **kwargs)

    def run(self, input_data: object, **kwargs: object) -> object:
        return asyncio.get_event_loop().run_until_complete(
            self._protected_invoke(input_data, **kwargs)
        )

    async def __call__(self, input_data: object, **kwargs: object) -> object:
        return await self._protected_invoke(input_data, **kwargs)

    # ------------------------------------------------------------------
    # Internals
    # ------------------------------------------------------------------

    async def _emit_event(
        self,
        event_type: str,
        tool_name: str | None,
        tool_params: dict | None,
        risk_score: int | None,
        reason: str | None,
        attack_type: str | None,
        decision: str | None,
    ) -> None:
        event = SentinelyEvent(
            org_id=self._org_id,
            agent_id=self._agent_id,
            session_id=self._session_id,
            event_type=event_type,
            tool_name=tool_name,
            tool_params=tool_params,
            risk_score=risk_score,
            reason=reason,
            attack_type=attack_type,
            decision=decision,
            task_hash=self._task_hash,
            action_chain=list(self._action_chain),
            timestamp=datetime.now(timezone.utc).isoformat(),
        )
        await self._emitter.emit(event)

    async def _post_fingerprint(self) -> None:
        """Fire-and-forget: POST the current fingerprint snapshot to the API.

        Skipped in development mode (no API key configured) so the SDK works
        locally without a running backend.
        """
        if not self._emitter.api_key or self._emitter.env == "development":
            return

        try:
            payload = self._fingerprint.to_api_payload()
            payload["session_id"] = self._session_id
            async with httpx.AsyncClient() as client:
                await client.post(
                    f"{self._emitter.api_url}/api/v1/agents/{self._agent_id}/fingerprint",
                    json=payload,
                    headers={"Authorization": f"Bearer {self._emitter.api_key}"},
                    timeout=1.0,
                )
        except Exception as exc:
            # Never let fingerprint posting crash the agent.
            logger.warning(
                "sentinely: failed to post fingerprint for %s: %s",
                self._agent_id,
                exc,
            )

    # ------------------------------------------------------------------
    # Public accessors
    # ------------------------------------------------------------------

    @property
    def memory(self) -> MemoryFirewall:
        return self._memory_firewall

    @property
    def fingerprint(self) -> BehavioralFingerprint:
        """Current behavioral fingerprint (read-only snapshot for testing)."""
        return self._fingerprint

    def get_audit_log(self) -> list[SentinelyEvent]:
        return self._emitter.get_local_queue()
