"""
LockStock OpenAI Agents SDK Guardrails
---------------------------------------
Chain-based guardrail implementations for OpenAI Agents SDK.

NO SECRETS. ONLY CHAIN STATE.

The Guard daemon tracks chain state (hash, matrix, sequence).
When signing, Guard uses current_hash as the HMAC key.

Usage:
    from agents import Agent
    from lockstock_openai import LockStockGuardrail

    # Chain Authority Mode (recommended)
    guardrail = LockStockGuardrail.from_liberty("agent_abc123")

    agent = Agent(
        name="my-agent",
        instructions="You are a helpful assistant.",
        guardrails=[guardrail]
    )
"""

from typing import Any, Dict, Optional
from dataclasses import dataclass

from lockstock_core.types import map_tool_to_capability


@dataclass
class GuardrailResult:
    """Result of a guardrail check."""
    passed: bool
    reason: Optional[str] = None
    modified_content: Optional[Any] = None


class LockStockGuardrail:
    """
    Chain-based guardrail for OpenAI Agents SDK.

    NO SECRETS STORED. Chain state is the authentication.

    The Guard daemon:
    1. Tracks current chain state (hash, matrix, sequence)
    2. Uses current_hash as HMAC key for signatures
    3. Returns signatures (not secrets!)
    """

    name = "lockstock"
    description = "LockStock chain-based authorization guardrail"

    def __init__(
        self,
        agent_id: str,
        guard_socket: str = "/var/run/lockstock-guard/guard.sock",
        endpoint: str = "https://lockstock-api-i9kp.onrender.com",
        block_on_failure: bool = True
    ):
        """
        Initialize chain-based guardrail.

        Args:
            agent_id: The agent's unique identifier
            guard_socket: Path to Guard daemon Unix socket
            endpoint: LockStock API endpoint
            block_on_failure: If True, block execution on auth failure

        NO SECRET PARAMETER! Secrets don't exist in this system.
        """
        self.agent_id = agent_id
        self.guard_socket = guard_socket
        self.endpoint = endpoint
        self.block_on_failure = block_on_failure

        # Chain state (cached locally, synced from Guard)
        self.current_hash: Optional[str] = None
        self.current_sequence: int = 0
        self.current_matrix: Optional[Dict[str, int]] = None
        self._chain_initialized = False

        # Guard client (for chain operations)
        from lockstock_guard.client import LibertyClient
        self._guard = LibertyClient(socket_path=guard_socket)

        # HTTP client (for server communication)
        import httpx
        self._http = httpx.AsyncClient(timeout=30.0)

    @classmethod
    def from_liberty(
        cls,
        agent_id: str,
        socket_path: str = "/var/run/lockstock-guard/guard.sock",
        endpoint: str = "https://lockstock-api-i9kp.onrender.com",
        block_on_failure: bool = True
    ):
        """
        Create guardrail using Chain Authority Mode.

        NO SECRET RETRIEVAL. NO SECRET STORAGE.

        The Guard daemon manages chain state. We just track it locally.

        Args:
            agent_id: Agent identifier (from provisioning)
            socket_path: Path to Guard daemon Unix socket
            endpoint: LockStock API endpoint
            block_on_failure: Whether to block on auth failure

        Returns:
            LockStockGuardrail instance using chain-based auth

        Example:
            >>> guardrail = LockStockGuardrail.from_liberty("agent_abc123")
            >>> agent = Agent(guardrails=[guardrail])
        """
        return cls(
            agent_id=agent_id,
            guard_socket=socket_path,
            endpoint=endpoint,
            block_on_failure=block_on_failure
        )

    async def _ensure_chain_initialized(self):
        """
        Sync chain state from Guard daemon on first use.

        Guard fetches current state from server and caches it.
        We sync that cached state to our local tracking.
        """
        if self._chain_initialized:
            return

        try:
            # Ask Guard to sync from server
            sync_result = self._guard.sync_chain(self.agent_id)

            if not sync_result.get("success"):
                raise Exception(f"Chain sync failed: {sync_result}")

            # Update local chain state cache
            self.current_hash = sync_result.get("current_hash")
            self.current_sequence = sync_result.get("sequence", 0)
            self.current_matrix = sync_result.get("state_matrix")
            self._chain_initialized = True

        except Exception as e:
            raise Exception(
                f"Failed to initialize chain state for {self.agent_id}. "
                f"Ensure Guard daemon is running and agent is provisioned. "
                f"Error: {e}"
            )

    async def validate(
        self,
        agent: Any,
        message: Dict[str, Any]
    ) -> GuardrailResult:
        """
        Validate a message using chain-based authentication.

        Flow:
        1. Check if action requires authorization
        2. Ask Guard to sign (Guard uses current_hash as HMAC key)
        3. Send Guard's signature to server for verification
        4. Update local chain state if verified

        NO SECRETS are involved. Only chain state.

        Args:
            agent: The agent instance
            message: The message to validate

        Returns:
            GuardrailResult indicating if validation passed
        """
        # Check for tool calls in assistant messages
        if message.get("role") == "assistant" and "tool_calls" in message:
            for tool_call in message.get("tool_calls", []):
                if "function" not in tool_call:
                    continue

                function_name = tool_call["function"].get("name", "")

                # Map tool to capability
                capability = map_tool_to_capability(function_name)

                # Ensure chain state is synced
                await self._ensure_chain_initialized()

                try:
                    # ========================================================
                    # CHAIN AUTHORITY MODE
                    # ========================================================
                    # Guard daemon:
                    # 1. Looks up current chain state for this agent
                    # 2. Calculates: new_matrix = current_matrix × generator[task]
                    # 3. Signs with: HMAC(current_hash, message)
                    #    ↑ current_hash IS the key! No separate secret!
                    # 4. Calculates: new_hash = SHA256(...)
                    # 5. Returns: {signature, new_hash, new_matrix, ...}
                    # ========================================================

                    sign_result = self._guard.sign_and_advance(
                        agent_id=self.agent_id,
                        task=capability,
                        parent_hash=self.current_hash
                    )

                    # Guard returned:
                    # - signature: HMAC computed using current_hash as key
                    # - new_hash: Next chain head
                    # - new_sequence: Incremented sequence
                    # - new_matrix: After applying generator
                    # - timestamp: Unix timestamp used in signature

                except Exception as e:
                    return GuardrailResult(
                        passed=False,
                        reason=f"Guard signing failed: {e}"
                    )

                # Send Guard's signature to LockStock server for verification
                try:
                    verify_result = await self._verify_with_server(
                        task=capability,
                        signature=sign_result["signature"],
                        parent_hash=self.current_hash,
                        state_matrix=sign_result["new_matrix"],
                        timestamp=sign_result["timestamp"],
                        sequence=sign_result["new_sequence"]
                    )

                    if not verify_result["authorized"]:
                        if self.block_on_failure:
                            return GuardrailResult(
                                passed=False,
                                reason=f"LockStock DENIED: {function_name} requires {capability}. "
                                       f"Reason: {verify_result.get('reason', 'Unknown')}"
                            )

                    # Update local chain state cache
                    self.current_hash = sign_result["new_hash"]
                    self.current_sequence = sign_result["new_sequence"]
                    self.current_matrix = sign_result["new_matrix"]

                except Exception as e:
                    return GuardrailResult(
                        passed=False,
                        reason=f"Server verification failed: {e}"
                    )

        return GuardrailResult(passed=True)

    async def _verify_with_server(
        self,
        task: str,
        signature: str,
        parent_hash: str,
        state_matrix: Dict[str, int],
        timestamp: int,
        sequence: int
    ) -> Dict[str, Any]:
        """
        Send Guard-signed request to LockStock server.

        The signature was created by Guard using current_hash as HMAC key.
        Server will verify by:
        1. Looking up its current_hash for this agent
        2. Recalculating HMAC using that hash
        3. Comparing to signature we send

        Args:
            task: Capability being verified
            signature: HMAC signature from Guard (signed with current_hash)
            parent_hash: Parent hash in chain
            state_matrix: New state matrix
            timestamp: Timestamp used in signature
            sequence: Sequence number

        Returns:
            Dict with authorized (bool) and reason (str)
        """
        payload = {
            "client_id": self.agent_id,
            "task": task.lower(),
            "parent_hash": parent_hash,
            "state_matrix": state_matrix,
            "signature": signature,  # ← From Guard, signed with current_hash
            "timestamp": timestamp,
            "sequence": sequence
        }

        response = await self._http.post(
            f"{self.endpoint}/verify",
            json=payload
        )

        data = response.json()

        return {
            "authorized": data.get("status") == "accepted",
            "reason": data.get("reason"),
            "state_hash": data.get("state_hash")
        }

    async def close(self):
        """Close connections."""
        await self._http.aclose()
        self._guard.close()


class LockStockInputGuardrail(LockStockGuardrail):
    """
    Input guardrail using chain-based authentication.
    """

    name = "lockstock_input"
    description = "LockStock chain-based input validation guardrail"

    async def validate_input(
        self,
        agent: Any,
        user_input: str,
        context: Optional[Dict[str, Any]] = None
    ) -> GuardrailResult:
        """Validate user input."""
        # Input validation - pass through for now
        # Capability checks happen at tool use
        return GuardrailResult(passed=True)


class LockStockOutputGuardrail(LockStockGuardrail):
    """
    Output guardrail using chain-based authentication.
    """

    name = "lockstock_output"
    description = "LockStock chain-based output validation guardrail"

    async def validate_output(
        self,
        agent: Any,
        output: str,
        context: Optional[Dict[str, Any]] = None
    ) -> GuardrailResult:
        """Validate agent output."""
        # Output validation - pass through for now
        # Could log to audit trail here
        return GuardrailResult(passed=True)
