import requests
import json
from typing import Optional, Dict, Any, Union

class AgentShield:
    def __init__(self, api_key: str, proxy_url: str = "https://<YOUR_API_ID>.execute-api.us-east-2.amazonaws.com/dev"):
        """
        Initialize the AgentShield Client.
        
        :param api_key: Your sk_... key from the dashboard.
        :param proxy_url: The base URL of your deployed AgentShield instance.
        """
        self.api_key = api_key
        self.proxy_url = proxy_url.rstrip('/')
        self.session = requests.Session()
        self.session.headers.update({
            "x-agentshield-key": self.api_key,
            "Content-Type": "application/json"
        })

    def protect(
        self, 
        agent_id: str, 
        payload: Union[str, Dict[str, Any]], 
        method: str = "POST", 
        path: str = "/"
    ) -> Dict[str, Any]:
        """
        Send a payload through the AgentShield Proxy for inspection.
        
        :param agent_id: The ID of the agent making the request (e.g., "support-bot-1")
        :param payload: The body content (JSON dict or string) to inspect/redact.
        :param method: The HTTP method the agent intends to use (default: POST).
        :param path: The target path the agent is accessing (default: /).
        :return: A dictionary containing success status, redacted data, or block reason.
        """
        url = f"{self.proxy_url}/proxy{path}"
        
        # Ensure headers are set correctly for this request
        headers = {"X-Agent-ID": agent_id}
        
        try:
            # If payload is a string, send it as a raw body, otherwise as JSON
            if isinstance(payload, str):
                response = self.session.request(method, url, data=payload, headers=headers)
            else:
                response = self.session.request(method, url, json=payload, headers=headers)
            
            # Handle 403 Forbidden (Blocked by Firewall)
            if response.status_code == 403:
                return {
                    "success": False,
                    "blocked": True,
                    "error": "BLOCKED_BY_FIREWALL",
                    "details": response.json()
                }
                
            # Handle 200 OK (Allowed & Redacted)
            if response.status_code == 200:
                data = response.json()
                return {
                    "success": True,
                    "blocked": False,
                    "redacted": data.get('redacted', False),
                    "original_data": payload,
                    "safe_data": data.get('redacted_body', data.get('body')),
                    "metadata": {
                        "plan": data.get('plan'),
                        "ai_protection_enabled": data.get('ai_protection')
                    }
                }
            
            # Handle other errors (401, 500, etc.)
            return {
                "success": False,
                "blocked": False,
                "error": f"HTTP_ERROR_{response.status_code}",
                "details": response.text
            }

        except Exception as e:
            return {
                "success": False, 
                "blocked": False, 
                "error": "CONNECTION_ERROR", 
                "details": str(e)
            }