"""Kevros Governance Client."""

from __future__ import annotations

import httpx
from typing import Any, Dict, List, Optional

from .models import (
    AttestResponse,
    BindIntentResponse,
    BundleResponse,
    GatewayHealth,
    IntentSource,
    IntentType,
    VerifyOutcomeResponse,
    VerifyResponse,
)

DEFAULT_BASE_URL = "https://governance.taskhawktech.com"
_TIMEOUT = 30.0


class GovernanceError(Exception):
    """Raised when the governance gateway returns an error."""

    def __init__(self, status_code: int, detail: str):
        self.status_code = status_code
        self.detail = detail
        super().__init__(f"Request failed (HTTP {status_code})")


class GovernanceClient:
    """Client for the Kevros A2A Governance Gateway."""

    def __init__(
        self,
        api_key: str,
        base_url: str = DEFAULT_BASE_URL,
        timeout: float = _TIMEOUT,
    ):
        self._api_key = api_key
        self._base_url = base_url.rstrip("/")
        self._timeout = timeout
        self._headers = {
            "X-API-Key": api_key,
            "Content-Type": "application/json",
        }
        self._sync_client: Optional[httpx.Client] = None
        self._async_client: Optional[httpx.AsyncClient] = None

    def __enter__(self) -> GovernanceClient:
        self._sync_client = httpx.Client(
            base_url=self._base_url,
            headers=self._headers,
            timeout=self._timeout,
        )
        return self

    def __exit__(self, *args: Any) -> None:
        if self._sync_client:
            self._sync_client.close()
            self._sync_client = None

    async def __aenter__(self) -> GovernanceClient:
        self._async_client = httpx.AsyncClient(
            base_url=self._base_url,
            headers=self._headers,
            timeout=self._timeout,
        )
        return self

    async def __aexit__(self, *args: Any) -> None:
        if self._async_client:
            await self._async_client.aclose()
            self._async_client = None

    def _get_sync(self) -> httpx.Client:
        if self._sync_client is None:
            self._sync_client = httpx.Client(
                base_url=self._base_url,
                headers=self._headers,
                timeout=self._timeout,
            )
        return self._sync_client

    def _get_async(self) -> httpx.AsyncClient:
        if self._async_client is None:
            self._async_client = httpx.AsyncClient(
                base_url=self._base_url,
                headers=self._headers,
                timeout=self._timeout,
            )
        return self._async_client

    def _handle_response(self, resp: httpx.Response) -> Dict[str, Any]:
        if resp.status_code >= 400:
            raise GovernanceError(resp.status_code, "")
        return resp.json()


    def verify(
        self,
        action_type: str,
        action_payload: Dict[str, Any],
        agent_id: str,
        policy_context: Optional[Dict[str, Any]] = None,
        idempotency_key: Optional[str] = None,
    ) -> VerifyResponse:
        """Verify an action against policy bounds before executing it."""
        body: Dict[str, Any] = {
            "action_type": action_type,
            "action_payload": action_payload,
            "agent_id": agent_id,
        }
        if policy_context is not None:
            body["policy_context"] = policy_context
        if idempotency_key is not None:
            body["idempotency_key"] = idempotency_key

        resp = self._get_sync().post("/governance/verify", json=body)
        return VerifyResponse.from_dict(self._handle_response(resp))

    async def averify(
        self,
        action_type: str,
        action_payload: Dict[str, Any],
        agent_id: str,
        policy_context: Optional[Dict[str, Any]] = None,
        idempotency_key: Optional[str] = None,
    ) -> VerifyResponse:
        """Async version of verify()."""
        body: Dict[str, Any] = {
            "action_type": action_type,
            "action_payload": action_payload,
            "agent_id": agent_id,
        }
        if policy_context is not None:
            body["policy_context"] = policy_context
        if idempotency_key is not None:
            body["idempotency_key"] = idempotency_key

        resp = await self._get_async().post("/governance/verify", json=body)
        return VerifyResponse.from_dict(self._handle_response(resp))


    def attest(
        self,
        agent_id: str,
        action_description: str,
        action_payload: Dict[str, Any],
        context: Optional[Dict[str, Any]] = None,
        prior_attestation_hash: Optional[str] = None,
    ) -> AttestResponse:
        """Create a provenance record for an action."""
        body: Dict[str, Any] = {
            "agent_id": agent_id,
            "action_description": action_description,
            "action_payload": action_payload,
        }
        if context is not None:
            body["context"] = context
        if prior_attestation_hash is not None:
            body["prior_attestation_hash"] = prior_attestation_hash

        resp = self._get_sync().post("/governance/attest", json=body)
        return AttestResponse.from_dict(self._handle_response(resp))

    async def aattest(
        self,
        agent_id: str,
        action_description: str,
        action_payload: Dict[str, Any],
        context: Optional[Dict[str, Any]] = None,
        prior_attestation_hash: Optional[str] = None,
    ) -> AttestResponse:
        """Async version of attest()."""
        body: Dict[str, Any] = {
            "agent_id": agent_id,
            "action_description": action_description,
            "action_payload": action_payload,
        }
        if context is not None:
            body["context"] = context
        if prior_attestation_hash is not None:
            body["prior_attestation_hash"] = prior_attestation_hash

        resp = await self._get_async().post("/governance/attest", json=body)
        return AttestResponse.from_dict(self._handle_response(resp))


    def bind(
        self,
        agent_id: str,
        intent_type: IntentType,
        intent_description: str,
        command_payload: Dict[str, Any],
        intent_source: IntentSource = IntentSource.AI_PLANNER,
        goal_state: Optional[Dict[str, Any]] = None,
        max_duration_ms: Optional[float] = None,
        parent_intent_id: Optional[str] = None,
    ) -> BindIntentResponse:
        """Bind an intent to a command."""
        body: Dict[str, Any] = {
            "agent_id": agent_id,
            "intent_type": intent_type.value,
            "intent_description": intent_description,
            "command_payload": command_payload,
            "intent_source": intent_source.value,
        }
        if goal_state is not None:
            body["goal_state"] = goal_state
        if max_duration_ms is not None:
            body["max_duration_ms"] = max_duration_ms
        if parent_intent_id is not None:
            body["parent_intent_id"] = parent_intent_id

        resp = self._get_sync().post("/governance/bind", json=body)
        return BindIntentResponse.from_dict(self._handle_response(resp))

    async def abind(
        self,
        agent_id: str,
        intent_type: IntentType,
        intent_description: str,
        command_payload: Dict[str, Any],
        intent_source: IntentSource = IntentSource.AI_PLANNER,
        goal_state: Optional[Dict[str, Any]] = None,
        max_duration_ms: Optional[float] = None,
        parent_intent_id: Optional[str] = None,
    ) -> BindIntentResponse:
        """Async version of bind()."""
        body: Dict[str, Any] = {
            "agent_id": agent_id,
            "intent_type": intent_type.value,
            "intent_description": intent_description,
            "command_payload": command_payload,
            "intent_source": intent_source.value,
        }
        if goal_state is not None:
            body["goal_state"] = goal_state
        if max_duration_ms is not None:
            body["max_duration_ms"] = max_duration_ms
        if parent_intent_id is not None:
            body["parent_intent_id"] = parent_intent_id

        resp = await self._get_async().post("/governance/bind", json=body)
        return BindIntentResponse.from_dict(self._handle_response(resp))


    def verify_outcome(
        self,
        agent_id: str,
        intent_id: str,
        binding_id: str,
        actual_state: Dict[str, Any],
        tolerance: float = 0.1,
    ) -> VerifyOutcomeResponse:
        """Verify that an action achieved its declared intent."""
        body = {
            "agent_id": agent_id,
            "intent_id": intent_id,
            "binding_id": binding_id,
            "actual_state": actual_state,
            "tolerance": tolerance,
        }
        resp = self._get_sync().post("/governance/verify-outcome", json=body)
        return VerifyOutcomeResponse.from_dict(self._handle_response(resp))

    async def averify_outcome(
        self,
        agent_id: str,
        intent_id: str,
        binding_id: str,
        actual_state: Dict[str, Any],
        tolerance: float = 0.1,
    ) -> VerifyOutcomeResponse:
        """Async version of verify_outcome()."""
        body = {
            "agent_id": agent_id,
            "intent_id": intent_id,
            "binding_id": binding_id,
            "actual_state": actual_state,
            "tolerance": tolerance,
        }
        resp = await self._get_async().post("/governance/verify-outcome", json=body)
        return VerifyOutcomeResponse.from_dict(self._handle_response(resp))


    def bundle(
        self,
        agent_id: str,
        time_range_start: Optional[str] = None,
        time_range_end: Optional[str] = None,
        include_intent_chains: bool = True,
        include_pqc_signatures: bool = True,
        include_verification_instructions: bool = True,
    ) -> BundleResponse:
        """Generate a compliance evidence bundle."""
        body: Dict[str, Any] = {
            "agent_id": agent_id,
            "include_intent_chains": include_intent_chains,
            "include_pqc_signatures": include_pqc_signatures,
            "include_verification_instructions": include_verification_instructions,
        }
        if time_range_start is not None:
            body["time_range_start"] = time_range_start
        if time_range_end is not None:
            body["time_range_end"] = time_range_end

        resp = self._get_sync().post("/governance/bundle", json=body)
        return BundleResponse.from_dict(self._handle_response(resp))

    async def abundle(
        self,
        agent_id: str,
        time_range_start: Optional[str] = None,
        time_range_end: Optional[str] = None,
        include_intent_chains: bool = True,
        include_pqc_signatures: bool = True,
        include_verification_instructions: bool = True,
    ) -> BundleResponse:
        """Async version of bundle()."""
        body: Dict[str, Any] = {
            "agent_id": agent_id,
            "include_intent_chains": include_intent_chains,
            "include_pqc_signatures": include_pqc_signatures,
            "include_verification_instructions": include_verification_instructions,
        }
        if time_range_start is not None:
            body["time_range_start"] = time_range_start
        if time_range_end is not None:
            body["time_range_end"] = time_range_end

        resp = await self._get_async().post("/governance/bundle", json=body)
        return BundleResponse.from_dict(self._handle_response(resp))


    def health(self) -> GatewayHealth:
        """Check governance gateway health."""
        resp = self._get_sync().get("/governance/health")
        return GatewayHealth.from_dict(self._handle_response(resp))

    async def ahealth(self) -> GatewayHealth:
        """Async version of health()."""
        resp = await self._get_async().get("/governance/health")
        return GatewayHealth.from_dict(self._handle_response(resp))
