"""
MCP Client Wrapper for Tupl SDK.

Provides a Pythonic interface for configuring agent rules and enforcement
via the Tupl MCP server.
"""

import json
from typing import Optional, Dict, Any, Callable
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client

from .agent import enforcement_agent as _enforcement_agent


class TuplAgentWrapper:
    """
    Client wrapper for integrating Tupl rule-based enforcement with user agents.

    Features:
    - Configure agent rule families via MCP
    - Discover available rule families
    - Wrap agents with enforcement
    - Access telemetry and compiled rules

    Example:
        ```python
        from tupl.mcp_client import TuplAgentWrapper

        # Initialize wrapper
        tupl = TuplAgentWrapper(tenant_id="my-team")

        # Configure agent rules
        config = await tupl.configure_agent(
            agent_id="cs-agent",
            owner="ops-team",
            rule_families={
                "tool_whitelist": {
                    "enabled": True,
                    "params": {
                        "allowed_tool_ids": ["search", "update"],
                        "action": "DENY"
                    }
                },
                "output_pii": {
                    "enabled": True,
                    "params": {"action": "REDACT"}
                }
            }
        )

        # Wrap agent with enforcement
        secure_agent = await tupl.wrap_agent(
            agent=my_langgraph_agent,
            agent_id="cs-agent",
            enforcement_mode="block"
        )

        # Use agent normally
        result = secure_agent.invoke({"messages": [...]})
        ```
    """

    def __init__(
        self,
        tenant_id: str,
        mcp_server_command: list[str] = None,
        control_plane_url: str = "http://localhost:8000",
        management_plane_url: str = "http://localhost:8000",
    ):
        """
        Initialize Tupl wrapper with MCP connection.

        Args:
            tenant_id: Your tenant identifier
            mcp_server_command: Command to start MCP server
            control_plane_url: Policy Control Plane URL
            management_plane_url: Management Plane URL (for enforcement)
        """
        self.tenant_id = tenant_id
        self.control_plane_url = control_plane_url
        self.management_plane_url = management_plane_url
        self.mcp_command = mcp_server_command or ["uv", "run", "mcp-tupl-server"]

        self._session: Optional[ClientSession] = None
        self._read = None
        self._write = None

    async def _ensure_session(self):
        """Ensure MCP session is connected."""
        if self._session is None:
            server_params = StdioServerParameters(
                command=self.mcp_command[0],
                args=self.mcp_command[1:],
                env={
                    "CONTROL_PLANE_URL": self.control_plane_url,
                    "MANAGEMENT_PLANE_URL": self.management_plane_url
                }
            )

            self._read, self._write = await stdio_client(server_params)
            self._session = ClientSession(self._read, self._write)
            await self._session.initialize()

    async def configure_agent(
        self,
        agent_id: str,
        owner: str,
        rule_families: Dict[str, Any],
        description: str = None,
    ) -> Dict[str, Any]:
        """
        Configure rule families for an agent.

        Args:
            agent_id: Agent identifier
            owner: Owner/team name
            rule_families: Dictionary of rule family configurations
            description: Optional agent description

        Returns:
            Configuration response with compiled rules

        Example:
            ```python
            config = await tupl.configure_agent(
                agent_id="cs-agent",
                owner="ops-team",
                rule_families={
                    "tool_whitelist": {
                        "enabled": True,
                        "params": {
                            "allowed_tool_ids": ["search_customer", "update_email"],
                            "action": "DENY"
                        }
                    },
                    "output_pii": {
                        "enabled": True,
                        "params": {
                            "action": "REDACT",
                            "pii_types": ["SSN", "CREDIT_CARD"]
                        }
                    }
                }
            )

            print(f"Configured {config['rule_count']} rules")
            ```
        """
        await self._ensure_session()

        result = await self._session.call_tool(
            "tupl_configure_agent_rules",
            arguments={
                "agent_id": agent_id,
                "owner": owner,
                "rule_families": rule_families,
                "description": description
            }
        )

        return json.loads(result.content[0].text)

    async def get_agent_rules(
        self,
        agent_id: str,
        include_compiled: bool = True
    ) -> Dict[str, Any]:
        """
        Get rule configuration for an agent.

        Args:
            agent_id: Agent identifier
            include_compiled: Include compiled RuleInstances

        Returns:
            Rule configuration
        """
        await self._ensure_session()

        result = await self._session.call_tool(
            "tupl_get_agent_rules",
            arguments={
                "agent_id": agent_id,
                "include_rules": include_compiled
            }
        )

        return json.loads(result.content[0].text)

    async def list_rule_families(
        self,
        layer: Optional[str] = None
    ) -> list[Dict[str, Any]]:
        """
        List available rule families.

        Args:
            layer: Optional layer filter (L0-L6)

        Returns:
            List of rule families with schemas

        Example:
            ```python
            families = await tupl.list_rule_families(layer="L4")

            for family in families["rule_families"]:
                print(f"{family['id']}: {family['description']}")
                print(f"  Parameters: {family['parameters']}")
            ```
        """
        await self._ensure_session()

        arguments = {}
        if layer:
            arguments["layer"] = layer

        result = await self._session.call_tool(
            "tupl_list_rule_families",
            arguments=arguments
        )

        return json.loads(result.content[0].text)

    async def get_rule_family_schema(
        self,
        family_id: str
    ) -> Dict[str, Any]:
        """
        Get detailed schema for a specific rule family.

        Args:
            family_id: Rule family identifier

        Returns:
            Complete rule family schema with parameter details

        Example:
            ```python
            schema = await tupl.get_rule_family_schema("tool_whitelist")

            print(f"Family: {schema['name']}")
            print(f"Layer: {schema['layer']}")
            print(f"Parameters:")
            for param, spec in schema['params_schema']['properties'].items():
                print(f"  - {param}: {spec['description']}")
            ```
        """
        await self._ensure_session()

        result = await self._session.call_tool(
            "tupl_get_rule_family_schema",
            arguments={"family_id": family_id}
        )

        return json.loads(result.content[0].text)

    async def wrap_agent(
        self,
        agent: Any,
        agent_id: str,
        boundary_id: str = "all",
        enforcement_mode: str = "audit",
        action_mapper: Optional[Callable] = None,
        resource_type_mapper: Optional[Callable] = None,
    ) -> Any:
        """
        Wrap a LangGraph agent with Tupl enforcement.

        Args:
            agent: Compiled LangGraph agent
            agent_id: Agent identifier (should match configured agent)
            boundary_id: Policy boundary to enforce
            enforcement_mode: "block" or "audit"
            action_mapper: Custom tool→action mapping
            resource_type_mapper: Custom resource type inference

        Returns:
            Secured agent (SecureGraphProxy instance)

        Example:
            ```python
            # First configure agent rules
            await tupl.configure_agent(
                agent_id="cs-agent",
                owner="ops-team",
                rule_families={...}
            )

            # Then wrap agent
            secure_agent = await tupl.wrap_agent(
                agent=my_agent,
                agent_id="cs-agent",
                enforcement_mode="block"
            )

            # Use normally - rules are enforced
            result = secure_agent.invoke({"messages": [...]})
            ```
        """
        await self._ensure_session()

        # Get enforcement config from MCP
        config_result = await self._session.call_tool(
            "tupl_wrap_agent",
            arguments={
                "agent_id": agent_id,
                "boundary_id": boundary_id,
                "enforcement_mode": enforcement_mode,
                "tenant_id": self.tenant_id
            }
        )

        config = json.loads(config_result.content[0].text)

        # Wrap agent using SDK
        secure_agent = _enforcement_agent(
            graph=agent,
            boundary_id=config["config"]["boundary_id"],
            tenant_id=config["config"]["tenant_id"],
            base_url=config["config"]["base_url"],
            action_mapper=action_mapper,
            resource_type_mapper=resource_type_mapper,
        )

        return secure_agent

    async def test_intent(
        self,
        action: str,
        actor_type: str,
        resource_type: str,
        **kwargs
    ) -> Dict[str, Any]:
        """
        Test an intent against boundaries.

        Args:
            action: Action type (read, write, delete, etc.)
            actor_type: Actor type (user, agent, service, llm)
            resource_type: Resource type (database, file, api)
            **kwargs: Additional intent parameters

        Returns:
            Decision result with similarities

        Example:
            ```python
            result = await tupl.test_intent(
                action="delete",
                actor_type="agent",
                resource_type="database",
                resource_name="customers",
                data={"pii": True, "sensitivity": ["internal"]},
                risk={"authn": "required"}
            )

            if result["decision"] == 0:
                print("Intent would be BLOCKED")
            else:
                print("Intent would be ALLOWED")
            print(f"Similarities: {result['slice_similarities']}")
            ```
        """
        await self._ensure_session()

        intent = {
            "action": action,
            "actor": {"type": actor_type, "id": kwargs.get("actor_id", "test")},
            "resource": {"type": resource_type, "name": kwargs.get("resource_name")},
            "data": kwargs.get("data", {"sensitivity": ["internal"], "pii": False}),
            "risk": kwargs.get("risk", {"authn": "required"}),
        }

        result = await self._session.call_tool(
            "tupl_test_intent",
            arguments={
                "intent_event": intent,
                "tenant_id": self.tenant_id,
            }
        )

        return json.loads(result.content[0].text)

    async def get_telemetry(
        self,
        limit: int = 100,
        tenant_id: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Retrieve enforcement telemetry.

        Args:
            limit: Max events to return
            tenant_id: Optional tenant filter

        Returns:
            Telemetry response with events

        Example:
            ```python
            telemetry = await tupl.get_telemetry(limit=50)

            print(f"Found {telemetry['total']} events")
            for event in telemetry['events']:
                print(f"  {event['intent_id']}: {event['decision']}")
            ```
        """
        await self._ensure_session()

        result = await self._session.call_tool(
            "tupl_get_telemetry",
            arguments={
                "limit": limit,
                "tenant_id": tenant_id or self.tenant_id,
            }
        )

        return json.loads(result.content[0].text)

    async def close(self):
        """Close MCP session."""
        if self._session:
            await self._session.close()
            self._session = None
