"""SecurityMixin — input/output validation, PII, guardrails, audit."""
from __future__ import annotations

from typing import Any
from uuid import uuid4

from logos_adk.errors import GuardrailError, InputValidationError, PIIDetectedError, PromptInjectionError
from logos_adk.types import Response, content_text


class SecurityMixin:
    """Mixin providing security methods for Agent."""

    def _resolve_pii_action(self, action: str, *, blocked: bool) -> str:
        if action == "selective":
            return "block" if blocked else "warn"
        return action

    def _audit(
        self,
        *,
        event: str,
        outcome: str,
        session_id: str | None,
        trace_id: str,
        request_id: str | None,
        attributes: dict[str, Any] | None = None,
    ) -> None:
        if not self._audit_logging_enabled:
            return
        self._audit_logger.log(
            event=event,
            outcome=outcome,
            agent_name=self.name,
            session_id=session_id,
            trace_id=trace_id,
            request_id=request_id,
            attributes=attributes,
        )

    async def _secure_prompt(
        self,
        prompt: str,
        *,
        session_id: str | None,
        trace_id: str,
        request_id: str | None,
    ) -> str:
        if self._input_validation_enabled:
            validation = await self._input_validator.validate(prompt)
            if not validation.passed:
                self._audit(
                    event="input_validation",
                    outcome="blocked",
                    session_id=session_id,
                    trace_id=trace_id,
                    request_id=request_id,
                    attributes=validation.metadata | {"message": validation.message},
                )
                raise InputValidationError(validation.message)

        if self._prompt_shield_enabled:
            shield = await self._prompt_shield.analyze(prompt)
            if shield.is_injection:
                self._audit(
                    event="prompt_injection",
                    outcome="blocked",
                    session_id=session_id,
                    trace_id=trace_id,
                    request_id=request_id,
                    attributes={"reason": shield.reason, "patterns": shield.matched_patterns},
                )
                raise PromptInjectionError("Prompt injection detected")

        if self._pii_detection_enabled:
            pii = await self._pii_detector.scan(prompt)
            pii_action = self._resolve_pii_action(self._input_pii_action, blocked=pii.blocked)
            if pii.has_pii and pii_action == "block":
                self._audit(
                    event="pii_detection",
                    outcome="blocked",
                    session_id=session_id,
                    trace_id=trace_id,
                    request_id=request_id,
                    attributes={"findings": [finding.__dict__ for finding in pii.findings]},
                )
                raise PIIDetectedError("Blocked PII detected in input")
            if pii.has_pii and pii_action == "warn":
                self._audit(
                    event="pii_detection",
                    outcome="warning",
                    session_id=session_id,
                    trace_id=trace_id,
                    request_id=request_id,
                    attributes={"findings": [finding.__dict__ for finding in pii.findings]},
                )
            else:
                self._audit(
                    event="input_validation",
                    outcome="allowed",
                    session_id=session_id,
                    trace_id=trace_id,
                    request_id=request_id,
                )
        return prompt

    async def _secure_output(
        self,
        response: Response,
        *,
        session_id: str | None,
        trace_id: str,
        request_id: str | None,
    ) -> Response:
        response_text = content_text(response.content, response.content_parts)
        if len(response_text) > self._output_max_length:
            self._audit(
                event="output_validation",
                outcome="blocked",
                session_id=session_id,
                trace_id=trace_id,
                request_id=request_id,
                attributes={"max_length": self._output_max_length, "length": len(response_text)},
            )
            raise GuardrailError(f"Output exceeds max length of {self._output_max_length} characters")
        if not self._pii_detection_enabled:
            return response
        pii = await self._pii_detector.scan(response_text)
        pii_action = self._resolve_pii_action(self._output_pii_action, blocked=pii.blocked)
        if pii.has_pii and pii_action == "block":
            self._audit(
                event="output_pii_detection",
                outcome="blocked",
                session_id=session_id,
                trace_id=trace_id,
                request_id=request_id,
                attributes={"findings": [finding.__dict__ for finding in pii.findings]},
            )
            raise PIIDetectedError("Blocked PII detected in output")
        if pii.has_pii and pii_action == "warn":
            self._audit(
                event="output_pii_detection",
                outcome="warning",
                session_id=session_id,
                trace_id=trace_id,
                request_id=request_id,
                attributes={"findings": [finding.__dict__ for finding in pii.findings]},
            )
        return response

    async def _check_guardrails(
        self,
        content: str,
        guardrails: list[Any],
        role: str,
        *,
        session_id: str | None = None,
        trace_id: str | None = None,
        request_id: str | None = None,
    ) -> str:
        """Run guardrails on content. Returns (possibly overridden) content.

        Raises GuardrailError on block action.
        """
        for guard in guardrails:
            result = await guard.check(content, role=role)
            if result.passed:
                continue

            self._audit(
                event="guardrail_triggered",
                outcome=result.action,
                session_id=session_id,
                trace_id=trace_id or uuid4().hex,
                request_id=request_id,
                attributes={
                    "guardrail": type(guard).__name__,
                    "role": role,
                    "message": result.message,
                    "action": result.action,
                },
            )

            if result.action == "block":
                raise GuardrailError(result.message, result=result)
            elif result.action == "override":
                content = result.override_content or content
            # "warn" — continue, span emitted by caller

        return content
