from __future__ import annotations
"""Async HTTP client for Vaikora API."""
import asyncio
from functools import wraps
from typing import Any, Callable, Optional, ParamSpec, TypeVar

import httpx

from vaikora.exceptions import NetworkError, TimeoutError, raise_for_status
from vaikora.models import (
    Action,
    ActionComplete,
    ActionResult,
    ActionSubmit,
    Agent,
    AgentCreate,
    AgentUpdate,
    Alert,
    PaginatedList,
    Policy,
    DataValidationRequest,
    DataValidationResult,
    DataValidationBatchResult,
)

P = ParamSpec("P")
T = TypeVar("T")

DEFAULT_BASE_URL = "https://api.vaikora.com"
DEFAULT_TIMEOUT = 30.0
DEFAULT_RETRY_COUNT = 3


class AgentsAPI:
    """API for managing agents."""

    def __init__(self, client: "VaikoraClient"):
        self._client = client

    async def register(
        self,
        name: str,
        agent_type: str = "autonomous",
        capabilities: Optional[list[str]] = None,
        metadata: Optional[dict[str, Any]] = None,
    ) -> Agent:
        """Register a new agent."""
        data = AgentCreate(
            name=name,
            agent_type=agent_type,  # type: ignore
            capabilities=capabilities or [],
            metadata=metadata or {},
        )
        response = await self._client._request("POST", "/agents", json=data.model_dump())
        return Agent.model_validate(response)

    async def get(self, agent_id: str) -> Agent:
        """Get an agent by ID."""
        response = await self._client._request("GET", f"/agents/{agent_id}")
        return Agent.model_validate(response)

    async def list(
        self,
        page: int = 1,
        page_size: int = 20,
        status: Optional[str] = None,
        agent_type: Optional[str] = None,
    ) -> PaginatedList:
        """List agents with pagination."""
        params: dict[str, Any] = {"page": page, "page_size": page_size}
        if status:
            params["status"] = status
        if agent_type:
            params["agent_type"] = agent_type
        response = await self._client._request("GET", "/agents", params=params)
        items = [Agent.model_validate(item) for item in response.get("items", [])]
        return PaginatedList(
            items=items,
            total=response.get("total", 0),
            page=response.get("page", page),
            page_size=response.get("page_size", page_size),
            has_more=response.get("has_more", False),
        )

    async def update(
        self,
        agent_id: str,
        name: Optional[str] = None,
        agent_type: Optional[str] = None,
        status: Optional[str] = None,
        capabilities: Optional[list[str]] = None,
        metadata: Optional[dict[str, Any]] = None,
    ) -> Agent:
        """Update an agent."""
        data = AgentUpdate(
            name=name,
            agent_type=agent_type,  # type: ignore
            status=status,  # type: ignore
            capabilities=capabilities,
            metadata=metadata,
        )
        # Remove None values
        update_data = {k: v for k, v in data.model_dump().items() if v is not None}
        response = await self._client._request("PATCH", f"/agents/{agent_id}", json=update_data)
        return Agent.model_validate(response)

    async def deactivate(self, agent_id: str) -> Agent:
        """Deactivate an agent."""
        return await self.update(agent_id, status="inactive")

    async def delete(self, agent_id: str) -> None:
        """Delete an agent."""
        await self._client._request("DELETE", f"/agents/{agent_id}")


class ActionsAPI:
    """API for managing actions."""

    def __init__(self, client: "VaikoraClient"):
        self._client = client

    async def submit(
        self,
        agent_id: str,
        action_type: str,
        resource: Optional[str] = None,
        payload: Optional[dict[str, Any]] = None,
        metadata: Optional[dict[str, Any]] = None,
    ) -> ActionResult:
        """Submit an action for policy evaluation."""
        data = ActionSubmit(
            agent_id=agent_id,  # type: ignore
            action_type=action_type,
            resource=resource,
            payload=payload or {},
            metadata=metadata or {},
        )
        response = await self._client._request("POST", "/actions/evaluate", json=data.model_dump())
        return ActionResult.model_validate(response)

    async def get(self, action_id: str) -> Action:
        """Get an action by ID."""
        response = await self._client._request("GET", f"/actions/{action_id}")
        return Action.model_validate(response)

    async def list(
        self,
        page: int = 1,
        page_size: int = 20,
        agent_id: Optional[str] = None,
        action_type: Optional[str] = None,
        status: Optional[str] = None,
        is_anomaly: Optional[bool] = None,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None,
    ) -> PaginatedList:
        """List actions with pagination and filters."""
        params: dict[str, Any] = {"page": page, "page_size": page_size}
        if agent_id:
            params["agent_id"] = agent_id
        if action_type:
            params["action_type"] = action_type
        if status:
            params["status"] = status
        if is_anomaly is not None:
            params["is_anomaly"] = is_anomaly
        if start_date:
            params["start_date"] = start_date
        if end_date:
            params["end_date"] = end_date
        response = await self._client._request("GET", "/actions", params=params)
        items = [Action.model_validate(item) for item in response.get("items", [])]
        return PaginatedList(
            items=items,
            total=response.get("total", 0),
            page=response.get("page", page),
            page_size=response.get("page_size", page_size),
            has_more=response.get("has_more", False),
        )

    async def complete(
        self,
        action_id: str,
        status: str,
        execution_time_ms: Optional[int] = None,
        error_message: Optional[str] = None,
        result_metadata: Optional[dict[str, Any]] = None,
    ) -> Action:
        """Mark an action as completed."""
        data = ActionComplete(
            status=status,  # type: ignore
            execution_time_ms=execution_time_ms,
            error_message=error_message,
            result_metadata=result_metadata or {},
        )
        response = await self._client._request(
            "POST", f"/actions/{action_id}/complete", json=data.model_dump()
        )
        return Action.model_validate(response)


class PoliciesAPI:
    """API for managing policies."""

    def __init__(self, client: "VaikoraClient"):
        self._client = client

    async def list(
        self,
        page: int = 1,
        page_size: int = 20,
        enabled: Optional[bool] = None,
        policy_type: Optional[str] = None,
    ) -> PaginatedList:
        """List policies with pagination."""
        params: dict[str, Any] = {"page": page, "page_size": page_size}
        if enabled is not None:
            params["enabled"] = enabled
        if policy_type:
            params["policy_type"] = policy_type
        response = await self._client._request("GET", "/policies", params=params)
        items = [Policy.model_validate(item) for item in response.get("items", [])]
        return PaginatedList(
            items=items,
            total=response.get("total", 0),
            page=response.get("page", page),
            page_size=response.get("page_size", page_size),
            has_more=response.get("has_more", False),
        )

    async def get(self, policy_id: str) -> Policy:
        """Get a policy by ID."""
        response = await self._client._request("GET", f"/policies/{policy_id}")
        return Policy.model_validate(response)


class AlertsAPI:
    """API for managing alerts."""

    def __init__(self, client: "VaikoraClient"):
        self._client = client

    async def list(
        self,
        page: int = 1,
        page_size: int = 20,
        status: Optional[str] = None,
        severity: Optional[str] = None,
        agent_id: Optional[str] = None,
    ) -> PaginatedList:
        """List alerts with pagination."""
        params: dict[str, Any] = {"page": page, "page_size": page_size}
        if status:
            params["status"] = status
        if severity:
            params["severity"] = severity
        if agent_id:
            params["agent_id"] = agent_id
        response = await self._client._request("GET", "/alerts", params=params)
        items = [Alert.model_validate(item) for item in response.get("items", [])]
        return PaginatedList(
            items=items,
            total=response.get("total", 0),
            page=response.get("page", page),
            page_size=response.get("page_size", page_size),
            has_more=response.get("has_more", False),
        )

    async def get(self, alert_id: str) -> Alert:
        """Get an alert by ID."""
        response = await self._client._request("GET", f"/alerts/{alert_id}")
        return Alert.model_validate(response)

    async def acknowledge(self, alert_id: str) -> Alert:
        """Acknowledge an alert."""
        response = await self._client._request("POST", f"/alerts/{alert_id}/acknowledge")
        return Alert.model_validate(response)

    async def resolve(self, alert_id: str, notes: Optional[str] = None) -> Alert:
        """Resolve an alert."""
        data = {"resolution_notes": notes} if notes else {}
        response = await self._client._request("POST", f"/alerts/{alert_id}/resolve", json=data)
        return Alert.model_validate(response)

    async def ignore(self, alert_id: str) -> Alert:
        """Ignore an alert."""
        response = await self._client._request("POST", f"/alerts/{alert_id}/ignore")
        return Alert.model_validate(response)


class ValidationAPI:
    """API for data validation."""

    def __init__(self, client: "VaikoraClient"):
        self._client = client

    async def validate(
        self,
        data: Any,
        check_pii: bool = True,
        check_anomalies: bool = True,
        check_toxicity: bool = True,
        auto_clean: bool = False,
        anomaly_threshold: float = 0.1,
        strict_mode: bool = False,
    ) -> DataValidationResult:
        """
        Validate data for quality, PII, anomalies, and toxicity.
        
        Args:
            data: The data to validate (dict, list, or string)
            check_pii: Enable PII detection
            check_anomalies: Enable anomaly detection
            check_toxicity: Enable toxicity detection
            auto_clean: Automatically clean data issues
            anomaly_threshold: Sensitivity for anomaly detection (0.0-1.0)
            strict_mode: Reject data on any issue found
            
        Returns:
            DataValidationResult with quality score and detailed findings
        """
        request = DataValidationRequest(
            data=data,
            check_pii=check_pii,
            check_anomalies=check_anomalies,
            check_toxicity=check_toxicity,
            auto_clean=auto_clean,
            anomaly_threshold=anomaly_threshold,
            strict_mode=strict_mode,
        )
        response = await self._client._request(
            "POST", "/validate", json=request.model_dump()
        )
        return DataValidationResult.model_validate(response)

    async def validate_batch(
        self,
        items: list[DataValidationRequest],
        fail_fast: bool = False,
    ) -> DataValidationBatchResult:
        """
        Validate multiple data items in batch.
        
        Args:
            items: List of validation requests
            fail_fast: Stop processing on first failure
            
        Returns:
            Batch result with individual results and overall score
        """
        response = await self._client._request(
            "POST",
            "/validate/batch",
            json={
                "items": [item.model_dump() for item in items],
                "fail_fast": fail_fast,
            },
        )
        return DataValidationBatchResult.model_validate(response)

    async def health(self) -> dict[str, Any]:
        """Check validation service health."""
        return await self._client._request("GET", "/validate/health")


class VaikoraClient:
    """Async client for the Vaikora API."""

    def __init__(
        self,
        api_key: str,
        base_url: str = DEFAULT_BASE_URL,
        timeout: float = DEFAULT_TIMEOUT,
        retry_count: int = DEFAULT_RETRY_COUNT,
        agent_id: Optional[str] = None,
    ):
        """
        Initialize the Vaikora client.
        
        Args:
            api_key: Your Vaikora API key
            base_url: Base URL for the Vaikora API
            timeout: Request timeout in seconds
            retry_count: Number of retries for failed requests
            agent_id: Default agent ID for action submissions
        """
        self._api_key = api_key
        self._base_url = base_url.rstrip("/")
        self._timeout = timeout
        self._retry_count = retry_count
        self._default_agent_id = agent_id
        
        self._client = httpx.AsyncClient(
            base_url=f"{self._base_url}/api/v1",
            headers={
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
                "User-Agent": "vaikora-python/0.1.0",
            },
            timeout=timeout,
        )
        
        # API namespaces
        self.agents = AgentsAPI(self)
        self.actions = ActionsAPI(self)
        self.policies = PoliciesAPI(self)
        self.alerts = AlertsAPI(self)
        self.validation = ValidationAPI(self)

    async def _request(
        self,
        method: str,
        path: str,
        params: Optional[dict[str, Any]] = None,
        json: Optional[dict[str, Any]] = None,
    ) -> dict[str, Any]:
        """Make an HTTP request with retry logic."""
        last_error: Optional[Exception] = None
        
        for attempt in range(self._retry_count):
            try:
                response = await self._client.request(
                    method=method,
                    url=path,
                    params=params,
                    json=json,
                )
                
                if response.status_code >= 400:
                    try:
                        data = response.json()
                    except Exception:
                        data = {"detail": response.text}
                    raise_for_status(response.status_code, data)
                
                if response.status_code == 204:
                    return {}
                
                return response.json()
                
            except httpx.TimeoutException as e:
                last_error = TimeoutError(f"Request timed out: {e}")
                if attempt < self._retry_count - 1:
                    await asyncio.sleep(2 ** attempt)  # Exponential backoff
                    
            except httpx.NetworkError as e:
                last_error = NetworkError(f"Network error: {e}")
                if attempt < self._retry_count - 1:
                    await asyncio.sleep(2 ** attempt)
        
        if last_error:
            raise last_error
        raise NetworkError("Unknown error occurred")

    async def close(self) -> None:
        """Close the HTTP client."""
        await self._client.aclose()

    async def __aenter__(self) -> "VaikoraClient":
        return self

    async def __aexit__(self, *args: Any) -> None:
        await self.close()

    def secure_action(
        self,
        action_type: str,
        resource: Optional[str] = None,
        agent_id: Optional[str] = None,
    ) -> Callable[[Callable[P, T]], Callable[P, T]]:
        """
        Decorator to wrap an action with Vaikora security evaluation.
        
        Usage:
            @client.secure_action(action_type="database_write", resource="users")
            async def write_user(user_id: int, data: dict):
                await db.write(user_id, data)
        """
        def decorator(func: Callable[P, T]) -> Callable[P, T]:
            @wraps(func)
            async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
                target_agent_id = agent_id or self._default_agent_id
                if not target_agent_id:
                    raise ValueError("agent_id must be provided either to decorator or client")
                
                # Build payload from function arguments
                payload = {
                    "args": args,
                    "kwargs": kwargs,
                }
                
                # Submit action for evaluation
                import time
                start_time = time.time()
                
                result = await self.actions.submit(
                    agent_id=target_agent_id,
                    action_type=action_type,
                    resource=resource,
                    payload=payload,
                    metadata={"function": func.__name__},
                )
                
                if not result.approved:
                    from vaikora.exceptions import PolicyViolationError
                    raise PolicyViolationError(
                        message=result.denial_reason or "Action denied by policy",
                        policy_id=str(result.policy_id) if result.policy_id else None,
                    )
                
                # Execute the action
                try:
                    output = await func(*args, **kwargs)  # type: ignore
                    
                    # Mark as executed
                    execution_time_ms = int((time.time() - start_time) * 1000)
                    await self.actions.complete(
                        action_id=str(result.action_id),
                        status="executed",
                        execution_time_ms=execution_time_ms,
                    )
                    
                    return output
                    
                except Exception as e:
                    # Mark as failed
                    execution_time_ms = int((time.time() - start_time) * 1000)
                    await self.actions.complete(
                        action_id=str(result.action_id),
                        status="failed",
                        execution_time_ms=execution_time_ms,
                        error_message=str(e),
                    )
                    raise
            
            return wrapper  # type: ignore
        return decorator

    async def validate_data(
        self,
        data: Any,
        check_pii: bool = True,
        check_anomalies: bool = True,
        check_toxicity: bool = True,
        auto_clean: bool = False,
        strict_mode: bool = False,
    ) -> DataValidationResult:
        """
        Convenience method to validate data.
        
        This is a shortcut for client.validation.validate().
        
        Args:
            data: The data to validate (dict, list, or string)
            check_pii: Enable PII detection (default: True)
            check_anomalies: Enable anomaly detection (default: True)
            check_toxicity: Enable toxicity detection (default: True)
            auto_clean: Automatically clean data issues (default: False)
            strict_mode: Reject data on any issue found (default: False)
            
        Returns:
            DataValidationResult with quality score and detailed findings
            
        Example:
            result = await client.validate_data({"email": "test@example.com"})
            if result.is_valid:
                print(f"Quality score: {result.quality_score}")
            else:
                print(f"Issues found: {result.pii_count} PII, {result.anomaly_count} anomalies")
        """
        return await self.validation.validate(
            data=data,
            check_pii=check_pii,
            check_anomalies=check_anomalies,
            check_toxicity=check_toxicity,
            auto_clean=auto_clean,
            strict_mode=strict_mode,
        )
