from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import httpx

DEFAULT_BASE_URL = "https://memory-api.locaith.com"
SDK_VERSION = "0.1.0"


class LocaithMemoryError(Exception):
    """Base exception for Locaith Memory SDK errors."""


class LocaithAuthError(LocaithMemoryError):
    """Raised when the SDK cannot authenticate against the MaaS API."""


class LocaithQuotaExceededError(LocaithMemoryError):
    """Raised when the authenticated tenant exceeds its configured quota."""


class LocaithAPIError(LocaithMemoryError):
    """Raised for non-auth, non-quota API failures."""


@dataclass(slots=True)
class _ClientConfig:
    base_url: str
    api_key: str
    timeout: float

    @property
    def headers(self) -> dict[str, str]:
        return {
            "Accept": "application/json",
            "Content-Type": "application/json",
            "X-API-Key": self.api_key,
            "User-Agent": f"locaith-memory-sdk/{SDK_VERSION}",
        }


def _normalize_base_url(base_url: str | None) -> str:
    return (base_url or DEFAULT_BASE_URL).strip().rstrip("/")


def _validate_api_key(api_key: str) -> str:
    cleaned = api_key.strip()
    if not cleaned:
        raise ValueError("`api_key` is required.")
    return cleaned


def _build_metadata(metadata: dict[str, Any] | None = None, **extras: Any) -> dict[str, Any]:
    payload = dict(metadata or {})
    for key, value in extras.items():
        if value is not None and key not in payload:
            payload[key] = value
    return payload


class LocaithMemoryClient:
    """Synchronous developer SDK for the Locaith Memory Platform."""

    def __init__(self, api_key: str, base_url: str | None = None, timeout: float = 30.0):
        self.config = _ClientConfig(
            base_url=_normalize_base_url(base_url),
            api_key=_validate_api_key(api_key),
            timeout=timeout,
        )
        self._client = httpx.Client(
            base_url=self.config.base_url,
            timeout=self.config.timeout,
            headers=self.config.headers,
        )

    def ingest(
        self,
        *,
        user_id: str,
        role: str,
        content: str,
        metadata: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        response = self._client.post(
            "/api/v1/memory/ingest",
            json={
                "user_id": user_id,
                "role": role,
                "content": content,
                "metadata": metadata or {},
            },
        )
        return self._parse_json(response)

    def retrieve(self, *, user_id: str, query: str) -> dict[str, Any]:
        response = self._client.get(
            "/api/v1/memory/retrieve",
            params={"user_id": user_id, "query": query},
        )
        return self._parse_json(response)

    def import_history(self, *, user_id: str, raw_dump: str, source_ai: str = "chatgpt") -> dict[str, Any]:
        response = self._client.post(
            "/api/v1/memory/import_history",
            json={
                "user_id": user_id,
                "raw_dump": raw_dump,
                "source_ai": source_ai,
            },
        )
        return self._parse_json(response)

    def consolidate(self, *, user_id: str | None = None, limit: int | None = None) -> dict[str, Any]:
        payload: dict[str, Any] = {}
        if user_id:
            payload["user_id"] = user_id
        if limit is not None:
            payload["limit"] = limit
        response = self._client.post("/api/v1/memory/consolidate", json=payload or {})
        return self._parse_json(response)

    def implant(
        self,
        user_id: str,
        context: str,
        *,
        role: str = "user",
        auto_consolidate: bool = True,
        metadata: dict[str, Any] | None = None,
        source: str = "python_sdk",
        consolidate_limit: int | None = None,
    ) -> dict[str, Any]:
        """Write fresh context into working memory and optionally consolidate it immediately."""

        ingest_result = self.ingest(
            user_id=user_id,
            role=role,
            content=context,
            metadata=_build_metadata(
                metadata,
                channel="python_sdk",
                source=source,
                sdk_method="implant",
                content_type="text",
            ),
        )
        consolidation_result: dict[str, Any] | None = None
        if auto_consolidate:
            consolidation_result = self.consolidate(user_id=user_id, limit=consolidate_limit)
        return {
            "ok": True,
            "user_id": user_id,
            "ingest": ingest_result,
            "consolidation": consolidation_result,
        }

    def recall(self, query: str, user_id: str) -> str:
        """Return the raw XML/text memory context for prompt injection."""

        payload = self.retrieve(user_id=user_id, query=query)
        return str(payload.get("memory_context", ""))

    def close(self) -> None:
        self._client.close()

    def __enter__(self) -> "LocaithMemoryClient":
        return self

    def __exit__(self, *_args: object) -> None:
        self.close()

    @staticmethod
    def _parse_json(response: httpx.Response) -> dict[str, Any]:
        _raise_for_status(response)
        payload = response.json()
        if isinstance(payload, dict):
            return payload
        raise LocaithAPIError("Expected a JSON object from the Locaith Memory API.")


class AsyncLocaithMemoryClient:
    """Async developer SDK for the Locaith Memory Platform."""

    def __init__(self, api_key: str, base_url: str | None = None, timeout: float = 30.0):
        self.config = _ClientConfig(
            base_url=_normalize_base_url(base_url),
            api_key=_validate_api_key(api_key),
            timeout=timeout,
        )
        self._client = httpx.AsyncClient(
            base_url=self.config.base_url,
            timeout=self.config.timeout,
            headers=self.config.headers,
        )

    async def ingest(
        self,
        *,
        user_id: str,
        role: str,
        content: str,
        metadata: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        response = await self._client.post(
            "/api/v1/memory/ingest",
            json={
                "user_id": user_id,
                "role": role,
                "content": content,
                "metadata": metadata or {},
            },
        )
        return self._parse_json(response)

    async def retrieve(self, *, user_id: str, query: str) -> dict[str, Any]:
        response = await self._client.get(
            "/api/v1/memory/retrieve",
            params={"user_id": user_id, "query": query},
        )
        return self._parse_json(response)

    async def import_history(self, *, user_id: str, raw_dump: str, source_ai: str = "chatgpt") -> dict[str, Any]:
        response = await self._client.post(
            "/api/v1/memory/import_history",
            json={
                "user_id": user_id,
                "raw_dump": raw_dump,
                "source_ai": source_ai,
            },
        )
        return self._parse_json(response)

    async def consolidate(self, *, user_id: str | None = None, limit: int | None = None) -> dict[str, Any]:
        payload: dict[str, Any] = {}
        if user_id:
            payload["user_id"] = user_id
        if limit is not None:
            payload["limit"] = limit
        response = await self._client.post("/api/v1/memory/consolidate", json=payload or {})
        return self._parse_json(response)

    async def implant(
        self,
        user_id: str,
        context: str,
        *,
        role: str = "user",
        auto_consolidate: bool = True,
        metadata: dict[str, Any] | None = None,
        source: str = "python_sdk",
        consolidate_limit: int | None = None,
    ) -> dict[str, Any]:
        ingest_result = await self.ingest(
            user_id=user_id,
            role=role,
            content=context,
            metadata=_build_metadata(
                metadata,
                channel="python_sdk",
                source=source,
                sdk_method="implant",
                content_type="text",
            ),
        )
        consolidation_result: dict[str, Any] | None = None
        if auto_consolidate:
            consolidation_result = await self.consolidate(user_id=user_id, limit=consolidate_limit)
        return {
            "ok": True,
            "user_id": user_id,
            "ingest": ingest_result,
            "consolidation": consolidation_result,
        }

    async def recall(self, query: str, user_id: str) -> str:
        payload = await self.retrieve(user_id=user_id, query=query)
        return str(payload.get("memory_context", ""))

    async def aclose(self) -> None:
        await self._client.aclose()

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

    async def __aexit__(self, *_args: object) -> None:
        await self.aclose()

    @staticmethod
    def _parse_json(response: httpx.Response) -> dict[str, Any]:
        _raise_for_status(response)
        payload = response.json()
        if isinstance(payload, dict):
            return payload
        raise LocaithAPIError("Expected a JSON object from the Locaith Memory API.")


def _raise_for_status(response: httpx.Response) -> None:
    if response.status_code == 429:
        raise LocaithQuotaExceededError(_extract_error_message(response))
    if response.status_code in {401, 403}:
        raise LocaithAuthError(_extract_error_message(response))
    try:
        response.raise_for_status()
    except httpx.HTTPStatusError as exc:
        raise LocaithAPIError(_extract_error_message(response)) from exc


def _extract_error_message(response: httpx.Response) -> str:
    try:
        payload = response.json()
        if isinstance(payload, dict) and payload.get("detail"):
            return str(payload["detail"])
    except ValueError:
        pass
    return response.text.strip() or f"Locaith Memory API returned HTTP {response.status_code}."
