"""
MCP ClientSession wrapper for scandar-guard.

Drop-in replacement for MCP ClientSession that inspects all tool calls
and results. The most important integration — MCP tool results are the
primary vector for runtime prompt injection attacks.
"""
from __future__ import annotations

from typing import Any, Optional

from ..audit import AuditLogger
from ..config import GuardConfig
from ..guard import ScandarBlockedError
from ..session import SessionManager
from ..tool_inspector import ToolCallInspector


class GuardedMCPSession:
    """
    Transparent wrapper for MCP ClientSession with runtime security monitoring.

    Inspects all tool calls before execution and all tool results after.
    This catches prompt injection embedded in external content (web pages,
    files, APIs) before it reaches the language model.

    Usage::

        from mcp import ClientSession
        from scandar_guard.integrations.mcp import GuardedMCPSession

        async with GuardedMCPSession(ClientSession(read, write)) as session:
            result = await session.call_tool("read_file", {"path": "doc.txt"})
    """

    def __init__(
        self,
        session: Any,
        config: Optional[GuardConfig] = None,
    ) -> None:
        self._session = session
        self._config = config or GuardConfig()
        self._tool_inspector = ToolCallInspector()
        self._session_manager = SessionManager(session_id=self._config.session_id)
        self._logger = AuditLogger(log_file=self._config.log_file)

    async def call_tool(
        self,
        name: str,
        arguments: Optional[dict[str, Any]] = None,
    ) -> Any:
        """
        Call an MCP tool with full security inspection.

        1. Inspects tool call arguments for PII and injection patterns
        2. Executes the tool call on the real session
        3. Inspects the tool result for prompt injection — the critical step
        """
        args = arguments or {}

        # Step 1: Inspect arguments before sending
        call_result = self._tool_inspector.inspect_tool_call(
            name, args, self._session_manager.session_id
        )
        if call_result.findings:
            self._logger.log_event(
                "mcp_tool_call_inspected",
                self._session_manager.session_id,
                call_result.findings,
                {"tool_name": name, "arg_count": len(args)},
            )
        if self._config.mode == "block":
            for f in call_result.findings:
                if f.severity in self._config.block_on:
                    raise ScandarBlockedError(f)

        # Step 2: Call the real MCP session
        result = await self._session.call_tool(name, arguments)

        # Step 3: Inspect the result — catches injected instructions in external content
        result_content = self._extract_result_content(result)
        result_inspection = self._tool_inspector.inspect_tool_result(
            name, result_content, self._session_manager.session_id
        )
        if result_inspection.findings:
            self._logger.log_event(
                "mcp_tool_result_inspected",
                self._session_manager.session_id,
                result_inspection.findings,
                {"tool_name": name, "has_critical": not result_inspection.passed},
            )
        if self._config.mode == "block":
            for f in result_inspection.findings:
                if f.severity in self._config.block_on:
                    raise ScandarBlockedError(f)

        return result

    def _extract_result_content(self, result: Any) -> Any:
        """Extract inspectable content from MCP tool result."""
        if result is None:
            return ""
        # MCP SDK returns CallToolResult with .content list
        if hasattr(result, "content"):
            content = result.content
            if isinstance(content, list):
                parts: list[str] = []
                for item in content:
                    if hasattr(item, "text"):
                        parts.append(str(item.text))
                    elif isinstance(item, dict) and "text" in item:
                        parts.append(str(item["text"]))
                return "\n".join(parts)
            return str(content)
        return result

    async def __aenter__(self) -> "GuardedMCPSession":
        if hasattr(self._session, "__aenter__"):
            await self._session.__aenter__()
        return self

    async def __aexit__(self, *args: Any) -> None:
        if hasattr(self._session, "__aexit__"):
            await self._session.__aexit__(*args)

    def __getattr__(self, name: str) -> Any:
        """Proxy all other methods to the underlying session."""
        return getattr(self._session, name)
