"""
Kronvex Async Python SDK v0.3.0
Same interface as the sync client but with async/await.

Usage:
    from kronvex.async_client import AsyncKronvex

    async with AsyncKronvex(api_key="kx_live_...", agent_id="<uuid>") as kx:
        await kx.remember("User prefers dark mode")
        results = await kx.recall("user preferences")
        ctx = await kx.inject_context("What theme?")
"""

from __future__ import annotations

import os
from typing import Any, Optional

try:
    import httpx as _httpx
except ImportError:
    _httpx = None  # type: ignore

from kronvex.client import (
    Memory, RecallResult, RecallResponse, InjectContextResponse,
    KronvexError, AuthError, RateLimitError, NotFoundError,
    BASE_URL,
)


class AsyncKronvex:
    """
    Async Kronvex client — drop-in replacement for Kronvex with async/await.

    Args:
        api_key:  Your Kronvex API key. Reads KRONVEX_API_KEY env var if not provided.
        agent_id: Default agent UUID.
        base_url: Override API base URL.
        timeout:  Request timeout in seconds (default: 10).
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        agent_id: Optional[str] = None,
        base_url: str = BASE_URL,
        timeout: float = 10.0,
    ):
        if _httpx is None:
            raise ImportError("httpx is required for AsyncKronvex. pip install httpx")
        self.api_key = api_key or os.getenv("KRONVEX_API_KEY")
        if not self.api_key:
            raise AuthError("No API key. Pass api_key= or set KRONVEX_API_KEY.")
        self.agent_id = agent_id or os.getenv("KRONVEX_AGENT_ID")
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self._headers = {
            "X-API-Key": self.api_key,
            "Content-Type": "application/json",
            "User-Agent": "kronvex-python/0.3.0",
        }
        self._client: Optional[_httpx.AsyncClient] = None

    async def __aenter__(self):
        self._client = _httpx.AsyncClient(headers=self._headers, timeout=self.timeout)
        return self

    async def __aexit__(self, *_):
        if self._client:
            await self._client.aclose()

    # ── Core methods ──────────────────────────────────────────────────────────

    async def remember(
        self,
        content: str,
        *,
        memory_type: str = "episodic",
        session_id: Optional[str] = None,
        metadata: Optional[dict] = None,
        ttl_days: Optional[int] = None,
        pinned: bool = False,
        agent_id: Optional[str] = None,
    ) -> Memory:
        """Async: store a memory."""
        payload: dict[str, Any] = {"content": content, "memory_type": memory_type}
        if session_id: payload["session_id"] = session_id
        if metadata:   payload["metadata"] = metadata
        if ttl_days:   payload["ttl_days"] = ttl_days
        if pinned:     payload["pinned"] = True
        aid = self._resolve_agent(agent_id)
        data = await self._post(f"/agents/{aid}/remember", payload)
        return Memory._from_dict(data)

    async def recall(
        self,
        query: str,
        *,
        top_k: int = 5,
        threshold: float = 0.5,
        session_id: Optional[str] = None,
        memory_type: Optional[str] = None,
        agent_id: Optional[str] = None,
    ) -> RecallResponse:
        """Async: recall relevant memories."""
        payload: dict[str, Any] = {"query": query, "top_k": top_k, "threshold": threshold}
        if session_id:  payload["session_id"] = session_id
        if memory_type: payload["memory_type"] = memory_type
        aid = self._resolve_agent(agent_id)
        data = await self._post(f"/agents/{aid}/recall", payload)
        results = [
            RecallResult(
                memory=Memory._from_dict(r["memory"]),
                similarity=r.get("similarity", 0.0),
                confidence=r.get("confidence", 0.0),
            )
            for r in data.get("results", [])
        ]
        return RecallResponse(query=data.get("query", query), results=results, total_found=data.get("total_found", len(results)))

    async def inject_context(
        self,
        message: str,
        *,
        top_k: int = 5,
        threshold: float = 0.5,
        agent_id: Optional[str] = None,
    ) -> InjectContextResponse:
        """Async: inject context for a user message."""
        payload: dict[str, Any] = {"message": message, "top_k": top_k, "threshold": threshold}
        aid = self._resolve_agent(agent_id)
        data = await self._post(f"/agents/{aid}/inject-context", payload)
        results = [
            RecallResult(
                memory=Memory._from_dict(r["memory"]),
                similarity=r.get("similarity", 0.0),
                confidence=r.get("confidence", 0.0),
            )
            for r in data.get("memories", [])
        ]
        return InjectContextResponse(context_block=data.get("context_block", ""), memories_used=data.get("memories_used", 0), memories=results)

    async def create_agent(self, name: str, description: str = "") -> dict:
        return await self._post("/agents", {"name": name, "description": description})

    async def list_agents(self) -> list[dict]:
        return await self._get("/agents")

    async def list_sessions(self, agent_id: Optional[str] = None) -> list[dict]:
        aid = self._resolve_agent(agent_id)
        data = await self._get(f"/agents/{aid}/sessions")
        return data.get("sessions", [])

    async def delete_memory(self, memory_id: str, agent_id: Optional[str] = None) -> bool:
        aid = self._resolve_agent(agent_id)
        await self._delete(f"/agents/{aid}/memories/{memory_id}")
        return True

    # ── HTTP helpers ──────────────────────────────────────────────────────────

    def _resolve_agent(self, agent_id: Optional[str]) -> str:
        aid = agent_id or self.agent_id
        if not aid:
            raise ValueError("No agent_id. Set in AsyncKronvex(agent_id=...) or pass per call.")
        return aid

    def _get_client(self) -> _httpx.AsyncClient:
        if self._client:
            return self._client
        return _httpx.AsyncClient(headers=self._headers, timeout=self.timeout)

    async def _post(self, path: str, payload: dict) -> dict:
        async with self._get_client() as c:
            r = await c.post(self.base_url + path, json=payload, headers=self._headers)
        self._raise(r.status_code, r.text)
        return r.json()

    async def _get(self, path: str) -> Any:
        async with self._get_client() as c:
            r = await c.get(self.base_url + path, headers=self._headers)
        self._raise(r.status_code, r.text)
        return r.json()

    async def _delete(self, path: str) -> None:
        async with self._get_client() as c:
            r = await c.delete(self.base_url + path, headers=self._headers)
        self._raise(r.status_code, r.text)

    @staticmethod
    def _raise(status: int, body: str) -> None:
        if status < 400: return
        msg = f"HTTP {status}: {body[:200]}"
        if status == 401: raise AuthError(msg, status)
        if status == 404: raise NotFoundError(msg, status)
        if status == 429: raise RateLimitError(msg, status)
        raise KronvexError(msg, status)
