"""Aegis SDK client implementation"""

from typing import Any, Dict, List, Optional
import httpx
from .types import (
    ToolCall,
    ToolCallResult,
    Span,
    PolicyDecision,
    ApprovalRequest,
    TraceSession,
)
from .errors import (
    NetworkError,
    CapabilityBlockedError,
    ApprovalRequiredError,
)


class AegisClient:
    """Synchronous Aegis SDK client"""
    
    def __init__(
        self,
        base_url: str,
        api_key: Optional[str] = None,
        timeout: float = 30.0,
        retries: int = 3,
    ):
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self.retries = retries
        
        headers = {"Content-Type": "application/json"}
        if api_key:
            headers["Authorization"] = f"Bearer {api_key}"
        
        self.client = httpx.Client(
            base_url=self.base_url,
            timeout=timeout,
            headers=headers,
        )
    
    def __enter__(self):
        return self
    
    def __exit__(self, *args):
        self.client.close()
    
    def _handle_error(self, response: httpx.Response) -> None:
        """Handle HTTP error responses"""
        if response.status_code == 403:
            data = response.json()
            capability = data.get("capability")
            reason = data.get("reason")
            if capability:
                raise CapabilityBlockedError(capability, reason)
        
        if response.status_code == 202:
            data = response.json()
            approval_id = data.get("approvalId")
            capability = data.get("capability")
            if approval_id and capability:
                raise ApprovalRequiredError(approval_id, capability)
        
        raise NetworkError(
            response.json().get("message", response.text),
            response.status_code
        )
    
    def validate_tool_call(self, tool_call: ToolCall) -> ToolCallResult:
        """Validate a tool call against policy"""
        try:
            response = self.client.post(
                "/api/v1/validate",
                json=tool_call.model_dump()
            )
            response.raise_for_status()
            return ToolCallResult(**response.json())
        except ApprovalRequiredError as e:
            return ToolCallResult(
                allowed=False,
                requires_approval=True,
                approval_id=e.approval_id,
                reason=str(e),
            )
        except httpx.HTTPStatusError as e:
            self._handle_error(e.response)
            raise
    
    def create_span(self, span_data: Dict[str, Any]) -> Span:
        """Create a new span"""
        response = self.client.post("/api/v1/spans", json=span_data)
        response.raise_for_status()
        return Span(**response.json())
    
    def get_span(self, span_id: str) -> Span:
        """Get span by ID"""
        response = self.client.get(f"/api/v1/spans/{span_id}")
        response.raise_for_status()
        return Span(**response.json())
    
    def list_decisions(self, session_id: Optional[str] = None) -> List[PolicyDecision]:
        """List policy decisions"""
        params = {"sessionId": session_id} if session_id else {}
        response = self.client.get("/api/v1/decisions", params=params)
        response.raise_for_status()
        return [PolicyDecision(**d) for d in response.json()]
    
    def get_approval(self, approval_id: str) -> ApprovalRequest:
        """Get approval request"""
        response = self.client.get(f"/api/v1/approvals/{approval_id}")
        response.raise_for_status()
        return ApprovalRequest(**response.json())
    
    def list_approvals(self) -> List[ApprovalRequest]:
        """List pending approvals"""
        response = self.client.get("/api/v1/approvals")
        response.raise_for_status()
        return [ApprovalRequest(**a) for a in response.json()]
    
    def approve(self, approval_id: str) -> None:
        """Approve a request"""
        response = self.client.post(f"/api/v1/approvals/{approval_id}/approve")
        response.raise_for_status()
    
    def deny(self, approval_id: str, reason: str) -> None:
        """Deny a request"""
        response = self.client.post(
            f"/api/v1/approvals/{approval_id}/deny",
            json={"reason": reason}
        )
        response.raise_for_status()
    
    def get_trace(self, session_id: str) -> TraceSession:
        """Get trace session"""
        response = self.client.get(f"/api/v1/traces/{session_id}")
        response.raise_for_status()
        return TraceSession(**response.json())
    
    def verify_trace(self, bundle_data: bytes) -> Dict[str, Any]:
        """Verify trace bundle"""
        response = self.client.post(
            "/api/v1/traces/verify",
            content=bundle_data,
            headers={"Content-Type": "application/cbor"},
        )
        response.raise_for_status()
        return response.json()
    
    def health(self) -> Dict[str, str]:
        """Health check"""
        response = self.client.get("/health")
        response.raise_for_status()
        return response.json()


class AegisAsyncClient:
    """Asynchronous Aegis SDK client"""
    
    def __init__(
        self,
        base_url: str,
        api_key: Optional[str] = None,
        timeout: float = 30.0,
        retries: int = 3,
    ):
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self.retries = retries
        
        headers = {"Content-Type": "application/json"}
        if api_key:
            headers["Authorization"] = f"Bearer {api_key}"
        
        self.client = httpx.AsyncClient(
            base_url=self.base_url,
            timeout=timeout,
            headers=headers,
        )
    
    async def __aenter__(self):
        return self
    
    async def __aexit__(self, *args):
        await self.client.aclose()
    
    def _handle_error(self, response: httpx.Response) -> None:
        """Handle HTTP error responses"""
        if response.status_code == 403:
            data = response.json()
            capability = data.get("capability")
            reason = data.get("reason")
            if capability:
                raise CapabilityBlockedError(capability, reason)
        
        if response.status_code == 202:
            data = response.json()
            approval_id = data.get("approvalId")
            capability = data.get("capability")
            if approval_id and capability:
                raise ApprovalRequiredError(approval_id, capability)
        
        raise NetworkError(
            response.json().get("message", response.text),
            response.status_code
        )
    
    async def validate_tool_call(self, tool_call: ToolCall) -> ToolCallResult:
        """Validate a tool call against policy"""
        try:
            response = await self.client.post(
                "/api/v1/validate",
                json=tool_call.model_dump()
            )
            response.raise_for_status()
            return ToolCallResult(**response.json())
        except ApprovalRequiredError as e:
            return ToolCallResult(
                allowed=False,
                requires_approval=True,
                approval_id=e.approval_id,
                reason=str(e),
            )
        except httpx.HTTPStatusError as e:
            self._handle_error(e.response)
            raise
    
    async def create_span(self, span_data: Dict[str, Any]) -> Span:
        """Create a new span"""
        response = await self.client.post("/api/v1/spans", json=span_data)
        response.raise_for_status()
        return Span(**response.json())
    
    async def get_span(self, span_id: str) -> Span:
        """Get span by ID"""
        response = await self.client.get(f"/api/v1/spans/{span_id}")
        response.raise_for_status()
        return Span(**response.json())
    
    async def list_decisions(self, session_id: Optional[str] = None) -> List[PolicyDecision]:
        """List policy decisions"""
        params = {"sessionId": session_id} if session_id else {}
        response = await self.client.get("/api/v1/decisions", params=params)
        response.raise_for_status()
        return [PolicyDecision(**d) for d in response.json()]
    
    async def get_approval(self, approval_id: str) -> ApprovalRequest:
        """Get approval request"""
        response = await self.client.get(f"/api/v1/approvals/{approval_id}")
        response.raise_for_status()
        return ApprovalRequest(**response.json())
    
    async def list_approvals(self) -> List[ApprovalRequest]:
        """List pending approvals"""
        response = await self.client.get("/api/v1/approvals")
        response.raise_for_status()
        return [ApprovalRequest(**a) for a in response.json()]
    
    async def approve(self, approval_id: str) -> None:
        """Approve a request"""
        response = await self.client.post(f"/api/v1/approvals/{approval_id}/approve")
        response.raise_for_status()
    
    async def deny(self, approval_id: str, reason: str) -> None:
        """Deny a request"""
        response = await self.client.post(
            f"/api/v1/approvals/{approval_id}/deny",
            json={"reason": reason}
        )
        response.raise_for_status()
    
    async def get_trace(self, session_id: str) -> TraceSession:
        """Get trace session"""
        response = await self.client.get(f"/api/v1/traces/{session_id}")
        response.raise_for_status()
        return TraceSession(**response.json())
    
    async def verify_trace(self, bundle_data: bytes) -> Dict[str, Any]:
        """Verify trace bundle"""
        response = await self.client.post(
            "/api/v1/traces/verify",
            content=bundle_data,
            headers={"Content-Type": "application/cbor"},
        )
        response.raise_for_status()
        return response.json()
    
    async def health(self) -> Dict[str, str]:
        """Health check"""
        response = await self.client.get("/health")
        response.raise_for_status()
        return response.json()
