"""Typed async client for the Atomic REST API.

Atomic is the brain — this client is how every other DidiBrain component
(bootstrap, scraper, extractor, lint, didi_client) talks to it.

Design notes:
- One AtomicClient instance per long-running process. Reusable httpx.AsyncClient.
- Auth: Bearer token loaded from `settings.atomic_token`. If empty, public
  endpoints (/health, setup) still work.
- Errors: any non-2xx is wrapped in AtomicApiError with the response body.
- Pagination: helpers like `iter_atoms()` (added later) yield pages.

Only the endpoints we actually need are wrapped here. Add more as required;
do not pre-emptively wrap the full ~78-route surface.
"""

from __future__ import annotations

from collections.abc import AsyncIterator, Sequence
from contextlib import asynccontextmanager
from dataclasses import dataclass
from typing import Any

import httpx
from tenacity import (
    AsyncRetrying,
    retry_if_exception_type,
    stop_after_attempt,
    wait_exponential,
)

from shared.config import settings
from shared.logging import get_logger

log = get_logger(__name__)


class AtomicApiError(Exception):
    """Raised on any non-2xx response from atomic-server."""

    def __init__(self, message: str, *, status: int, body: str = "", url: str = ""):
        super().__init__(message)
        self.status = status
        self.body = body
        self.url = url


@dataclass(slots=True)
class AtomSummary:
    """Minimal atom view returned by list/search endpoints."""

    id: str
    content: str
    source_url: str | None
    embedding_status: str
    tagging_status: str
    created_at: str
    updated_at: str
    tags: list[dict[str, Any]]

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> AtomSummary:
        return cls(
            id=d["id"],
            content=d.get("content", ""),
            source_url=d.get("source_url"),
            embedding_status=d.get("embedding_status", "unknown"),
            tagging_status=d.get("tagging_status", "unknown"),
            created_at=d.get("created_at", ""),
            updated_at=d.get("updated_at", ""),
            tags=d.get("tags", []),
        )


@dataclass(slots=True)
class SearchHit:
    """One result from POST /api/search."""

    atom_id: str
    similarity: float
    matching_chunk_content: str | None
    snippet: str | None
    source_url: str | None
    tags: list[dict[str, Any]]

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> SearchHit:
        return cls(
            atom_id=d.get("id") or d.get("atom_id", ""),
            similarity=float(d.get("similarity_score") or d.get("similarity", 0.0)),
            matching_chunk_content=d.get("matching_chunk_content"),
            snippet=d.get("snippet"),
            source_url=d.get("source_url"),
            tags=d.get("tags", []),
        )


class AtomicClient:
    """Async REST client for atomic-server. Use as `async with AtomicClient() as a:`."""

    def __init__(
        self,
        *,
        base_url: str | None = None,
        token: str | None = None,
        timeout: float = 60.0,
        max_attempts: int = 3,
    ):
        self._base_url = (base_url or settings.atomic_url).rstrip("/")
        self._token = token if token is not None else settings.atomic_token
        self._max_attempts = max_attempts

        headers = {"Content-Type": "application/json"}
        if self._token:
            headers["Authorization"] = f"Bearer {self._token}"

        self._http = httpx.AsyncClient(
            base_url=self._base_url,
            headers=headers,
            timeout=httpx.Timeout(timeout, connect=10.0),
        )

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

    async def __aenter__(self) -> AtomicClient:
        return self

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

    @property
    def base_url(self) -> str:
        return self._base_url

    @property
    def has_token(self) -> bool:
        return bool(self._token)

    # ---------------------------------------------------------------- internals
    async def _request(
        self,
        method: str,
        path: str,
        *,
        json: Any = None,
        params: dict[str, Any] | None = None,
        retry_on_5xx: bool = True,
    ) -> Any:
        async for attempt in AsyncRetrying(
            stop=stop_after_attempt(self._max_attempts),
            wait=wait_exponential(multiplier=1, min=1, max=10),
            retry=retry_if_exception_type((httpx.HTTPError, httpx.TimeoutException)),
            reraise=True,
        ):
            with attempt:
                resp = await self._http.request(method, path, json=json, params=params)
                if resp.status_code >= 500 and retry_on_5xx:
                    raise httpx.HTTPError(f"5xx from {path}: {resp.status_code}")
                if resp.status_code >= 400:
                    raise AtomicApiError(
                        f"HTTP {resp.status_code} {method} {path}",
                        status=resp.status_code,
                        body=resp.text[:1000],
                        url=str(resp.url),
                    )
                if resp.status_code == 204 or not resp.content:
                    return None
                try:
                    return resp.json()
                except ValueError:
                    return resp.text
        raise AtomicApiError("retry loop exited without result", status=0)  # unreachable

    # ----------------------------------------------------------------- health
    async def health(self) -> dict[str, Any]:
        return await self._request("GET", "/health")

    async def setup_status(self) -> dict[str, Any]:
        """Whether the instance still needs initial token claim."""
        return await self._request("GET", "/api/setup/status")

    # ----------------------------------------------------------------- settings
    async def get_settings(self) -> dict[str, Any]:
        return await self._request("GET", "/api/settings")

    async def set_setting(self, key: str, value: str) -> Any:
        """Set a single setting. Use set_settings() for bulk."""
        return await self._request(
            "PUT", f"/api/settings/{key}", json={"value": value}
        )

    async def set_settings(self, items: dict[str, str]) -> dict[str, Any]:
        """Apply many settings sequentially. Returns map of key → response."""
        out: dict[str, Any] = {}
        for k, v in items.items():
            out[k] = await self.set_setting(k, v)
        return out

    # -------------------------------------------------------------------- tags
    async def list_tags(self, *, min_count: int = 0) -> list[dict[str, Any]]:
        return await self._request("GET", "/api/tags", params={"min_count": min_count})

    async def create_tag(
        self, name: str, *, parent_id: str | None = None
    ) -> dict[str, Any]:
        body: dict[str, Any] = {"name": name}
        if parent_id:
            body["parent_id"] = parent_id
        return await self._request("POST", "/api/tags", json=body)

    # ------------------------------------------------------------------- atoms
    async def create_atom(
        self,
        *,
        content: str,
        source_url: str | None = None,
        tag_ids: Sequence[str] | None = None,
        published_at: str | None = None,
    ) -> dict[str, Any]:
        body: dict[str, Any] = {"content": content, "tag_ids": list(tag_ids or [])}
        if source_url:
            body["source_url"] = source_url
        if published_at:
            body["published_at"] = published_at
        return await self._request("POST", "/api/atoms", json=body)

    async def get_atom(self, atom_id: str) -> dict[str, Any]:
        return await self._request("GET", f"/api/atoms/{atom_id}")

    async def delete_atom(self, atom_id: str) -> None:
        await self._request("DELETE", f"/api/atoms/{atom_id}")

    async def get_atom_by_source_url(self, source_url: str) -> dict[str, Any] | None:
        try:
            # Atomic's GetAtomBySourceUrlQuery uses `url` not `source_url`.
            return await self._request(
                "GET", "/api/atoms/by-source-url", params={"url": source_url}
            )
        except AtomicApiError as e:
            if e.status == 404:
                return None
            raise

    async def get_embedding_status(self, atom_id: str) -> dict[str, Any]:
        return await self._request("GET", f"/api/atoms/{atom_id}/embedding-status")

    async def list_atoms(
        self,
        *,
        limit: int = 50,
        offset: int = 0,
        tag_id: str | None = None,
    ) -> dict[str, Any]:
        params: dict[str, Any] = {"limit": limit, "offset": offset}
        if tag_id:
            params["tag_id"] = tag_id
        return await self._request("GET", "/api/atoms", params=params)

    # ------------------------------------------------------------------ search
    async def search(
        self,
        query: str,
        *,
        mode: str = "semantic",
        limit: int = 20,
        threshold: float | None = None,
    ) -> list[SearchHit]:
        body: dict[str, Any] = {"query": query, "mode": mode, "limit": limit}
        if threshold is not None:
            body["threshold"] = threshold
        result = await self._request("POST", "/api/search", json=body)
        # Atomic returns either a list directly or {"results": [...]}
        if isinstance(result, dict):
            items = result.get("results") or result.get("data") or []
        else:
            items = result or []
        return [SearchHit.from_dict(item) for item in items]

    async def find_similar(
        self, atom_id: str, *, threshold: float = 0.5, limit: int = 20
    ) -> list[SearchHit]:
        params = {"threshold": threshold, "limit": limit}
        result = await self._request(
            "GET", f"/api/atoms/{atom_id}/similar", params=params
        )
        items = result if isinstance(result, list) else result.get("results", [])
        return [SearchHit.from_dict(item) for item in items]

    # ------------------------------------------------------------- embeddings
    async def get_pipeline_status(self) -> dict[str, Any]:
        return await self._request("GET", "/api/embeddings/status")

    async def process_pending_embeddings(self) -> dict[str, Any]:
        return await self._request("POST", "/api/embeddings/process-pending")


@asynccontextmanager
async def atomic_client(
    *, token: str | None = None
) -> AsyncIterator[AtomicClient]:
    """`async with atomic_client() as a:` for short-lived scripts."""
    client = AtomicClient(token=token)
    try:
        yield client
    finally:
        await client.aclose()