"""Asynchronous client for Miruvor SDK."""
import asyncio
import logging
from typing import Any, Dict, List, Optional, Union

import httpx

from ._http import parse_json_or_raise
from .exceptions import MiruvorError, RateLimitError, ServerError
from .models import (
    IngestRequest,
    IngestResponse,
    RetrieveResponse,
    StoreRequest,
    StoreResponse,
)

logger = logging.getLogger("miruvor")


class AsyncMiruvorClient:
    """Asynchronous client for Zyra SNN Memory Database."""

    def __init__(
        self,
        base_url: str,
        api_key: str,
        token: Optional[str] = None,
        timeout: float = 30.0,
        max_retries: int = 3,
    ):
        """
        Initialize async Miruvor client.

        Args:
            base_url: API base URL
            api_key: X-API-Key for authentication
            token: Optional JWT token for tenant-scoped operations
            timeout: Request timeout in seconds
            max_retries: Max retry attempts for failed requests
        """
        self.base_url = base_url.rstrip("/")
        self.api_key = api_key
        self.token = token
        self.timeout = timeout
        self.max_retries = max_retries

        self._headers = {"X-API-Key": api_key}
        if token:
            self._headers["Authorization"] = f"Bearer {token}"

        self._client: Optional[httpx.AsyncClient] = None

    async def _request(
        self,
        method: str,
        path: str,
        *,
        json: Optional[Dict[str, Any]] = None,
        params: Optional[Dict[str, Any]] = None,
    ) -> Any:
        url = f"{self.base_url}{path}"
        for attempt in range(self.max_retries + 1):
            try:
                resp = await self.client.request(method, url, json=json, params=params)
                return parse_json_or_raise(resp)
            except (RateLimitError, ServerError) as e:
                if attempt >= self.max_retries:
                    logger.error(f"{method} {path} failed: {e.message} (status={e.status_code})")
                    raise
                delay = self._retry_delay_seconds(attempt, e)
                logger.warning(f"{method} {path} retrying in {delay:.2f}s after {e.status_code}")
                await asyncio.sleep(delay)
            except httpx.RequestError as e:
                if attempt >= self.max_retries:
                    raise MiruvorError(f"HTTP request failed: {e}") from e
                delay = self._retry_delay_seconds(attempt, None)
                logger.warning(f"{method} {path} network retry in {delay:.2f}s: {e}")
                await asyncio.sleep(delay)
            except MiruvorError as e:
                logger.error(f"{method} {path} failed: {e.message} (status={e.status_code})")
                raise
        raise MiruvorError(f"{method} {path} failed after retries")

    def _retry_delay_seconds(self, attempt: int, error: Optional[MiruvorError]) -> float:
        if isinstance(error, RateLimitError):
            return float(max(error.retry_after, 1))
        return min(8.0, 0.5 * (2**attempt))

    @property
    def client(self) -> httpx.AsyncClient:
        """Lazy client initialization to avoid event loop issues."""
        if self._client is None or self._client.is_closed:
            transport = httpx.AsyncHTTPTransport(retries=self.max_retries)
            self._client = httpx.AsyncClient(
                headers=self._headers,
                timeout=self.timeout,
                transport=transport,
            )
        return self._client

    async def health(self) -> Dict[str, Any]:
        logger.debug("Checking API health")
        return await self._request("GET", "/health")

    async def store(
        self,
        text: str,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> StoreResponse:
        clean_text = text.strip()
        if not clean_text:
            raise ValueError("Text cannot be empty")
        logger.debug(f"Storing memory: {clean_text[:50]}...")
        payload = StoreRequest(
            text=clean_text,
            tags=tags or [],
            metadata=metadata or {},
        ).model_dump()
        data = await self._request("POST", "/store", json=payload)
        logger.debug(f"Stored memory: {data.get('memory_id')}")
        return StoreResponse(**data)

    async def store_batch(
        self,
        memories: List[Dict[str, Any]],
        max_concurrent: int = 10,
    ) -> List[StoreResponse]:
        """
        Store multiple memories efficiently with concurrency control.

        Args:
            memories: List of dicts with 'text', optional 'tags', 'metadata'
            max_concurrent: Maximum concurrent requests

        Returns:
            List of StoreResponse objects
        """
        logger.debug(
            f"Storing {len(memories)} memories in batch (concurrency={max_concurrent})"
        )
        semaphore = asyncio.Semaphore(max_concurrent)

        async def store_one(memory: Dict[str, Any]) -> StoreResponse:
            async with semaphore:
                return await self.store(
                    text=memory["text"],
                    tags=memory.get("tags", []),
                    metadata=memory.get("metadata", {}),
                )

        return await asyncio.gather(*[store_one(m) for m in memories])

    async def ingest(
        self,
        content: Union[str, dict],
        priority: str = "normal",
        metadata: Optional[Dict[str, Any]] = None,
        enable_consolidation: bool = True,
        enable_background_linking: bool = True,
    ) -> IngestResponse:
        logger.debug(f"Ingesting content with priority={priority}")
        payload = IngestRequest(
            content=content,
            priority=priority,
            metadata=metadata or {},
            enable_consolidation=enable_consolidation,
            enable_background_linking=enable_background_linking,
        ).model_dump()
        data = await self._request("POST", "/ingest", json=payload)
        logger.debug(f"Ingested: message_id={data.get('message_id')}")
        return IngestResponse(**data)

    async def retrieve(
        self,
        query: str,
        top_k: int = 5,
        use_sparse: Optional[bool] = None,
    ) -> RetrieveResponse:
        logger.debug(f"Retrieving memories for query: {query[:50]}...")
        params: Dict[str, Any] = {"query": query, "top_k": top_k}
        if use_sparse is not None:
            params["use_sparse"] = use_sparse
        data = await self._request("POST", "/retrieve/pattern_completion", params=params)
        logger.debug(f"Retrieved {data.get('num_results', 0)} memories")
        return RetrieveResponse(**data)

    async def close(self) -> None:
        """Close the client."""
        if self._client and not self._client.is_closed:
            await self._client.aclose()
            self._client = None

    async def __aenter__(self) -> "AsyncMiruvorClient":
        """Async context manager entry."""
        transport = httpx.AsyncHTTPTransport(retries=self.max_retries)
        self._client = httpx.AsyncClient(
            headers=self._headers,
            timeout=self.timeout,
            transport=transport,
        )
        return self

    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        """Async context manager exit."""
        await self.close()
