import os
import json
import requests
from typing import Optional, Dict, Any, Tuple
from pydantic import BaseModel

from anchor8.identity_wallet import IdentityWallet


class GuardResponse(BaseModel):
    allowed: bool
    reason: Optional[str] = None
    modified_arguments: Optional[Dict[str, Any]] = None


class Anchor8Client:
    """
    Anchor8 SDK Client with Universal Agent Identity support.
    
    Provides:
    - Guard mode (Lane 2): Blocking security checks
    - Observer mode (Lane 1): Async telemetry logging
    - DID-Auth: Signed requests with agent identity
    """
    
    def __init__(
        self,
        api_key: Optional[str] = None,
        base_url: str = "http://localhost:8000/api/v1",
        private_key_jwk: Optional[Dict[str, str]] = None,
        did: Optional[str] = None,
        key_file: Optional[str] = None
    ):
        """
        Initialize the Anchor8 client.
        
        Args:
            api_key: API key for authentication
            base_url: Anchor8 API base URL
            private_key_jwk: Optional private key for DID-Auth
            did: Optional agent DID for DID-Auth
            key_file: Optional path to identity wallet file
        """
        self.api_key = api_key or os.getenv("ANCHOR8_API_KEY")
        self.base_url = base_url.rstrip("/")

        if not self.api_key:
            raise ValueError("API Key is required. Set ANCHOR8_API_KEY env var or pass in constructor.")

        self.session = requests.Session()
        self.session.headers.update({
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        })
        
        # Initialize identity wallet if credentials provided
        self._wallet: Optional[IdentityWallet] = None
        if private_key_jwk or did or key_file or os.getenv("ANCHOR8_AGENT_KEY"):
            self._wallet = IdentityWallet(
                private_key_jwk=private_key_jwk,
                did=did,
                key_file=key_file
            )
    
    @property
    def has_identity(self) -> bool:
        """Check if the client has a configured DID identity."""
        return self._wallet is not None and self._wallet.has_identity
    
    @property
    def did(self) -> Optional[str]:
        """Get the agent's DID if configured."""
        return self._wallet.did if self._wallet else None

    def check(
        self,
        agent_id: str,
        tool_name: str,
        arguments: Dict[str, Any],
        context: Optional[str] = None,
        sign_request: bool = True
    ) -> Tuple[bool, str, Optional[Dict[str, Any]]]:
        """
        Synchronous check with the Anchor8 Guard.
        
        If the client has a configured identity and sign_request=True,
        the request will be signed with the agent's private key.
        
        Args:
            agent_id: Agent identifier
            tool_name: Name of the tool being called
            arguments: Tool arguments to check
            context: Optional context string
            sign_request: Whether to sign the request with DID-Auth
        
        Returns:
            Tuple of (allowed, reason, modified_arguments)
        """
        payload = {
            "agent_id": agent_id,
            "tool_name": tool_name,
            "arguments": arguments,
            "context": context
        }
        
        # Prepare headers
        headers = {}
        
        # Add DID-Auth signature if available
        if sign_request and self._wallet and self._wallet.has_identity:
            auth_headers, _ = self._wallet.sign_request(payload)
            headers.update(auth_headers)

        try:
            response = self.session.post(
                f"{self.base_url}/guard/check",
                json=payload,
                headers=headers
            )
            response.raise_for_status()
            data = response.json()
            return data["allowed"], data.get("reason", ""), data.get("modified_arguments")
        except requests.exceptions.RequestException as e:
            # Fail-closed for security
            return False, f"Connection error: {str(e)}", None

    def log(
        self,
        agent_id: str,
        tool_name: str,
        metadata: Optional[Dict[str, Any]] = None
    ) -> None:
        """
        Async fire-and-forget log.
        """
        payload = {
            "agent_id": agent_id,
            "tool_name": tool_name,
            "metadata": metadata or {}
        }

        try:
            self.session.post(f"{self.base_url}/logs/", json=payload)
        except:
            # Logging should not break app flow
            pass
    
    def register_identity(
        self,
        owner_id: str,
        agent_name: str,
        domain: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Register this agent with the Anchor8 identity service.
        
        Creates a DID and stores the private key in the wallet.
        
        Args:
            owner_id: Enterprise owner identifier
            agent_name: Human-readable agent name
            domain: Optional domain for did:web
            
        Returns:
            Registration response with DID
        """
        payload = {
            "owner_id": owner_id,
            "agent_name": agent_name,
            "domain": domain
        }
        
        response = self.session.post(
            f"{self.base_url}/identity/register",
            json=payload
        )
        response.raise_for_status()
        data = response.json()
        
        # Store the identity in wallet
        if "private_key_jwk" in data and data["private_key_jwk"]:
            self._wallet = IdentityWallet(
                private_key_jwk=data["private_key_jwk"],
                did=data["did"]
            )
        
        return data
    
    def get_wallet(self) -> Optional[IdentityWallet]:
        """Get the identity wallet."""
        return self._wallet
    
    def save_identity(self, file_path: str) -> None:
        """
        Save the identity wallet to a file.
        
        Args:
            file_path: Path to save the wallet
        """
        if not self._wallet:
            raise ValueError("No identity configured")
        self._wallet.save_to_file(file_path)
    
    def load_identity(self, file_path: str) -> None:
        """
        Load identity wallet from a file.
        
        Args:
            file_path: Path to the wallet file
        """
        self._wallet = IdentityWallet(key_file=file_path)

