"""Small, version-pinned adapter over WeKnora's official REST API."""

from __future__ import annotations

import time
from collections.abc import Mapping, Sequence
from typing import Any
from urllib.parse import quote

import httpx


class WeKnoraError(RuntimeError):
    """Normalized upstream failure that is safe for DeerFlow to translate."""

    def __init__(self, message: str, *, status_code: int = 502) -> None:
        super().__init__(message)
        self.status_code = status_code


def _data(payload: Any) -> Any:
    if not isinstance(payload, dict):
        raise WeKnoraError("WeKnora returned an invalid response")
    if payload.get("success") is False:
        message = payload.get("message") or payload.get("error") or "WeKnora rejected the request"
        raise WeKnoraError(str(message), status_code=400)
    return payload.get("data", payload)


def _as_list(value: Any) -> list[dict[str, Any]]:
    if isinstance(value, list):
        return [dict(item) for item in value if isinstance(item, Mapping)]
    if isinstance(value, dict):
        for key in ("items", "list", "knowledge_bases", "knowledge", "records"):
            nested = value.get(key)
            if isinstance(nested, list):
                return [dict(item) for item in nested if isinstance(item, Mapping)]
    return []


class WeKnoraClient:
    """Async client for a trusted-network WeKnora endpoint.

    The caller may inject an ``httpx.AsyncClient`` for connection pooling and
    tests. When omitted, each public method creates and closes a short-lived client.
    If admin credentials are configured, DeerFlow authenticates to WeKnora with a
    cached bearer token; otherwise it keeps the legacy trusted-network anonymous
    mode for deployments that allow DeerFlow through a reverse proxy.
    """

    def __init__(
        self,
        *,
        base_url: str,
        timeout_seconds: float = 30.0,
        http_client: httpx.AsyncClient | None = None,
        admin_email: str = "",
        admin_password: str = "",
    ) -> None:
        self.base_url = base_url.rstrip("/")
        self.timeout_seconds = timeout_seconds
        self._http_client = http_client
        self.admin_email = admin_email.strip()
        self.admin_password = admin_password
        self._auth_token = ""
        self._tenant_id = ""
        self._auth_expires_at = 0.0
        if not self.base_url:
            raise ValueError("WeKnora base URL is required")

    @property
    def _has_admin_credentials(self) -> bool:
        return bool(self.admin_email and self.admin_password)

    @staticmethod
    def _extract_login_token(payload: Any) -> tuple[str, str]:
        if not isinstance(payload, dict):
            return "", ""
        source = payload.get("data") if isinstance(payload.get("data"), dict) else payload
        token = str(
            source.get("token")
            or source.get("access_token")
            or payload.get("token")
            or payload.get("access_token")
            or ""
        )
        tenant = (
            source.get("tenant")
            if isinstance(source.get("tenant"), dict)
            else source.get("active_tenant")
            if isinstance(source.get("active_tenant"), dict)
            else payload.get("tenant")
        )
        user = source.get("user") if isinstance(source.get("user"), dict) else payload.get("user")
        tenant_id = ""
        if isinstance(tenant, dict):
            tenant_id = str(tenant.get("id") or "")
        if not tenant_id and isinstance(user, dict):
            tenant_id = str(user.get("tenant_id") or "")
        return token, tenant_id

    async def _login(self, *, force_refresh: bool = False) -> None:
        if not self._has_admin_credentials:
            return
        now = time.time()
        if self._auth_token and not force_refresh and self._auth_expires_at > now + 60:
            return

        owns_client = self._http_client is None
        client = self._http_client or httpx.AsyncClient(
            timeout=httpx.Timeout(self.timeout_seconds),
            follow_redirects=False,
            trust_env=False,
        )
        try:
            response = await client.post(
                f"{self.base_url}/api/v1/auth/login",
                json={"email": self.admin_email, "password": self.admin_password},
                headers={"Accept": "application/json"},
            )
        except httpx.TimeoutException as exc:
            raise WeKnoraError("WeKnora admin login timed out", status_code=504) from exc
        except httpx.HTTPError as exc:
            raise WeKnoraError("WeKnora admin login is unavailable", status_code=502) from exc
        finally:
            if owns_client:
                await client.aclose()

        if response.status_code >= 400:
            raise WeKnoraError(f"WeKnora admin login failed ({response.status_code})", status_code=502)
        try:
            payload = response.json()
        except ValueError as exc:
            raise WeKnoraError("WeKnora admin login returned invalid JSON", status_code=502) from exc
        if isinstance(payload, dict) and payload.get("success") is False:
            raise WeKnoraError(str(payload.get("message") or "WeKnora admin login failed"), status_code=502)
        token, tenant_id = self._extract_login_token(payload)
        if not token:
            raise WeKnoraError("WeKnora admin login did not return a token", status_code=502)

        self._auth_token = token
        self._tenant_id = tenant_id
        # WeKnora login responses do not consistently expose expiry. Refresh
        # periodically and retry once if upstream says the token is stale.
        self._auth_expires_at = now + 25 * 60

    async def _headers(self, *, force_refresh: bool = False, accept: str = "application/json") -> dict[str, str]:
        headers = {"Accept": accept}
        if self._has_admin_credentials:
            await self._login(force_refresh=force_refresh)
            if self._auth_token:
                headers["Authorization"] = f"Bearer {self._auth_token}"
            if self._tenant_id:
                headers["X-Tenant-ID"] = self._tenant_id
        return headers

    async def _request(
        self,
        method: str,
        path: str,
        *,
        json: Any = None,
        params: Mapping[str, Any] | None = None,
        files: Any = None,
        data: Mapping[str, Any] | None = None,
    ) -> Any:
        owns_client = self._http_client is None
        client = self._http_client or httpx.AsyncClient(
            timeout=httpx.Timeout(self.timeout_seconds),
            follow_redirects=False,
            trust_env=False,
        )
        try:
            for attempt in range(2):
                response = await client.request(
                    method,
                    f"{self.base_url}{path}",
                    headers=await self._headers(force_refresh=attempt > 0),
                    json=json,
                    params=params,
                    files=files,
                    data=data,
                )
                if response.status_code not in {401, 403} or not self._has_admin_credentials or attempt > 0:
                    break
        except httpx.TimeoutException as exc:
            raise WeKnoraError("WeKnora request timed out", status_code=504) from exc
        except httpx.HTTPError as exc:
            raise WeKnoraError("WeKnora is unavailable", status_code=502) from exc
        finally:
            if owns_client:
                await client.aclose()
        if response.status_code == 404:
            raise WeKnoraError("WeKnora resource was not found", status_code=404)
        if response.status_code in {401, 403}:
            if self._has_admin_credentials:
                raise WeKnoraError("WeKnora rejected DeerFlow's authenticated internal request", status_code=502)
            raise WeKnoraError(
                "WeKnora rejected DeerFlow's anonymous internal request; allow DeerFlow in WeKnora or its reverse proxy",
                status_code=502,
            )
        if response.status_code >= 400:
            detail = ""
            try:
                body = response.json()
                if isinstance(body, dict):
                    detail = str(body.get("message") or body.get("error") or body.get("detail") or "")
            except ValueError:
                detail = ""
            raise WeKnoraError(detail or f"WeKnora returned HTTP {response.status_code}", status_code=502)
        if response.status_code == 204 or not response.content:
            return {"success": True}
        try:
            return response.json()
        except ValueError as exc:
            raise WeKnoraError("WeKnora returned invalid JSON") from exc

    async def health(self) -> dict[str, Any]:
        payload = await self._request("GET", "/health")
        return dict(payload) if isinstance(payload, dict) else {"status": "ok"}

    async def list_knowledge_bases(self) -> list[dict[str, Any]]:
        return _as_list(_data(await self._request("GET", "/api/v1/knowledge-bases")))

    async def get_knowledge_base(self, knowledge_base_id: str) -> dict[str, Any]:
        value = _data(await self._request("GET", f"/api/v1/knowledge-bases/{quote(knowledge_base_id, safe='')}"))
        if not isinstance(value, dict):
            raise WeKnoraError("WeKnora returned an invalid knowledge base")
        return dict(value)

    async def list_models(self) -> list[dict[str, Any]]:
        """List models from WeKnora; model credentials are never returned."""

        return _as_list(_data(await self._request("GET", "/api/v1/models")))

    async def get_default_embedding_model_id(self) -> str:
        models = [item for item in await self.list_models() if str(item.get("type") or "").strip().lower() == "embedding"]
        preferred = next((item for item in models if item.get("is_default") is True), None)
        selected = preferred or next(
            (item for item in models if str(item.get("status") or "active").lower() == "active"),
            None,
        )
        model_id = str((selected or {}).get("id") or "")
        if not model_id:
            raise WeKnoraError(
                "WeKnora has no available embedding model; configure its default model in WeKnora",
                status_code=400,
            )
        return model_id

    async def get_default_knowledge_qa_model_id(self) -> str:
        """Resolve WeKnora's own active synthesis model for Wiki indexing."""

        models = [item for item in await self.list_models() if str(item.get("type") or "").strip().lower() == "knowledgeqa"]
        preferred = next((item for item in models if item.get("is_default") is True), None)
        selected = preferred or next(
            (item for item in models if str(item.get("status") or "active").lower() == "active"),
            None,
        )
        model_id = str((selected or {}).get("id") or "")
        if not model_id:
            raise WeKnoraError(
                "WeKnora has no available KnowledgeQA model; configure one in WeKnora before creating a Wiki",
                status_code=400,
            )
        return model_id

    async def create_knowledge_base(
        self,
        *,
        name: str,
        description: str = "",
        kb_type: str = "document",
        embedding_model_id: str | None = None,
        wiki_enabled: bool = False,
    ) -> dict[str, Any]:
        model_id = embedding_model_id or await self.get_default_embedding_model_id()
        create_payload: dict[str, Any] = {
            "name": name,
            "description": description,
            "type": kb_type,
            "embedding_model_id": model_id,
        }
        if wiki_enabled:
            synthesis_model_id = await self.get_default_knowledge_qa_model_id()
            create_payload.update(
                {
                    "summary_model_id": synthesis_model_id,
                    "indexing_strategy": {
                        "vector_enabled": True,
                        "keyword_enabled": True,
                        "wiki_enabled": True,
                        "graph_enabled": False,
                    },
                    "wiki_config": {
                        "synthesis_model_id": synthesis_model_id,
                        "extraction_granularity": "standard",
                    },
                }
            )
        payload = await self._request(
            "POST",
            "/api/v1/knowledge-bases",
            json=create_payload,
        )
        value = _data(payload)
        if not isinstance(value, dict) or not value.get("id"):
            raise WeKnoraError("WeKnora did not return the created knowledge base")
        return dict(value)

    async def update_knowledge_base(self, knowledge_base_id: str, changes: Mapping[str, Any]) -> dict[str, Any]:
        payload = dict(changes)
        if not str(payload.get("name") or "").strip():
            current = await self.get_knowledge_base(knowledge_base_id)
            current_name = str(current.get("name") or "").strip()
            if not current_name:
                raise WeKnoraError("WeKnora knowledge base name is required for updates")
            payload["name"] = current_name
        value = _data(
            await self._request(
                "PUT",
                f"/api/v1/knowledge-bases/{quote(knowledge_base_id, safe='')}",
                json=payload,
            )
        )
        return dict(value) if isinstance(value, dict) else {}

    async def delete_knowledge_base(self, knowledge_base_id: str) -> None:
        await self._request("DELETE", f"/api/v1/knowledge-bases/{quote(knowledge_base_id, safe='')}")

    async def list_documents(self, knowledge_base_id: str, *, page: int = 1, page_size: int = 100) -> dict[str, Any]:
        payload = await self._request(
            "GET",
            f"/api/v1/knowledge-bases/{quote(knowledge_base_id, safe='')}/knowledge",
            params={"page": page, "page_size": page_size},
        )
        rows = _as_list(_data(payload))
        total = payload.get("total", len(rows)) if isinstance(payload, dict) else len(rows)
        return {"items": rows, "total": int(total or 0), "page": page, "page_size": page_size}

    async def upload_document(
        self,
        knowledge_base_id: str,
        *,
        filename: str,
        content: bytes,
        content_type: str = "application/octet-stream",
    ) -> dict[str, Any]:
        value = _data(
            await self._request(
                "POST",
                f"/api/v1/knowledge-bases/{quote(knowledge_base_id, safe='')}/knowledge/file",
                files={"file": (filename, content, content_type)},
            )
        )
        return dict(value) if isinstance(value, dict) else {}

    async def import_document_url(self, knowledge_base_id: str, *, url: str) -> dict[str, Any]:
        value = _data(
            await self._request(
                "POST",
                f"/api/v1/knowledge-bases/{quote(knowledge_base_id, safe='')}/knowledge/url",
                json={"url": url},
            )
        )
        return dict(value) if isinstance(value, dict) else {}

    async def create_manual_document(
        self,
        knowledge_base_id: str,
        *,
        title: str,
        content: str,
    ) -> dict[str, Any]:
        value = _data(
            await self._request(
                "POST",
                f"/api/v1/knowledge-bases/{quote(knowledge_base_id, safe='')}/knowledge/manual",
                json={"title": title, "content": content, "status": "publish"},
            )
        )
        return dict(value) if isinstance(value, dict) else {}

    async def delete_document(self, knowledge_id: str) -> None:
        await self._request("DELETE", f"/api/v1/knowledge/{quote(knowledge_id, safe='')}")

    async def get_document(self, knowledge_id: str) -> dict[str, Any]:
        value = _data(await self._request("GET", f"/api/v1/knowledge/{quote(knowledge_id, safe='')}"))
        if not isinstance(value, dict):
            raise WeKnoraError("WeKnora returned an invalid document")
        return dict(value)

    async def get_document_preview(self, knowledge_id: str) -> tuple[bytes, str]:
        """Return the original document preview bytes and its media type."""

        owns_client = self._http_client is None
        client = self._http_client or httpx.AsyncClient(
            timeout=httpx.Timeout(self.timeout_seconds),
            follow_redirects=True,
            trust_env=False,
        )
        try:
            for attempt in range(2):
                response = await client.request(
                    "GET",
                    f"{self.base_url}/api/v1/knowledge/{quote(knowledge_id, safe='')}/preview",
                    headers=await self._headers(force_refresh=attempt > 0, accept="*/*"),
                    follow_redirects=True,
                )
                if response.status_code not in {401, 403} or not self._has_admin_credentials or attempt > 0:
                    break
        except httpx.TimeoutException as exc:
            raise WeKnoraError("WeKnora preview request timed out", status_code=504) from exc
        except httpx.HTTPError as exc:
            raise WeKnoraError("WeKnora preview is unavailable", status_code=502) from exc
        finally:
            if owns_client:
                await client.aclose()
        if response.status_code == 404:
            raise WeKnoraError("WeKnora document preview was not found", status_code=404)
        if response.status_code in {401, 403}:
            if self._has_admin_credentials:
                raise WeKnoraError("WeKnora rejected DeerFlow's authenticated preview request", status_code=502)
            raise WeKnoraError(
                "WeKnora rejected DeerFlow's anonymous internal request; allow DeerFlow in WeKnora or its reverse proxy",
                status_code=502,
            )
        if response.status_code >= 400:
            raise WeKnoraError(f"WeKnora returned HTTP {response.status_code}", status_code=502)
        media_type = response.headers.get("content-type", "application/octet-stream").split(";", 1)[0].strip()
        return response.content, media_type or "application/octet-stream"

    async def get_knowledge_base_file(
        self,
        knowledge_base_id: str,
        file_path: str,
    ) -> tuple[bytes, str]:
        """Return one protected Wiki attachment using DeerFlow's WeKnora identity."""

        owns_client = self._http_client is None
        client = self._http_client or httpx.AsyncClient(
            timeout=httpx.Timeout(self.timeout_seconds),
            follow_redirects=True,
            trust_env=False,
        )
        try:
            for attempt in range(2):
                response = await client.request(
                    "GET",
                    f"{self.base_url}/api/v1/knowledge-bases/{quote(knowledge_base_id, safe='')}/files",
                    params={"file_path": file_path},
                    headers=await self._headers(force_refresh=attempt > 0, accept="*/*"),
                    follow_redirects=True,
                )
                if response.status_code not in {401, 403} or not self._has_admin_credentials or attempt > 0:
                    break
        except httpx.TimeoutException as exc:
            raise WeKnoraError("WeKnora Wiki image request timed out", status_code=504) from exc
        except httpx.HTTPError as exc:
            raise WeKnoraError("WeKnora Wiki image is unavailable", status_code=502) from exc
        finally:
            if owns_client:
                await client.aclose()
        if response.status_code == 404:
            raise WeKnoraError("WeKnora Wiki image was not found", status_code=404)
        if response.status_code in {401, 403}:
            if self._has_admin_credentials:
                raise WeKnoraError("WeKnora rejected DeerFlow's authenticated Wiki image request", status_code=502)
            raise WeKnoraError(
                "WeKnora rejected DeerFlow's anonymous internal request; allow DeerFlow in WeKnora or its reverse proxy",
                status_code=502,
            )
        if response.status_code >= 400:
            raise WeKnoraError(f"WeKnora returned HTTP {response.status_code}", status_code=502)
        media_type = response.headers.get("content-type", "application/octet-stream").split(";", 1)[0].strip()
        return response.content, media_type or "application/octet-stream"

    async def reprocess_document(self, knowledge_id: str) -> dict[str, Any]:
        value = _data(await self._request("POST", f"/api/v1/knowledge/{quote(knowledge_id, safe='')}/reparse"))
        return dict(value) if isinstance(value, dict) else {}

    async def list_chunks(self, knowledge_id: str, *, page: int = 1, page_size: int = 20) -> dict[str, Any]:
        payload = await self._request(
            "GET",
            f"/api/v1/chunks/{quote(knowledge_id, safe='')}",
            params={"page": page, "page_size": page_size},
        )
        rows = _as_list(_data(payload))
        total = payload.get("total", len(rows)) if isinstance(payload, dict) else len(rows)
        return {"items": rows, "total": int(total or 0), "page": page, "page_size": page_size}

    async def get_wiki_graph(self, knowledge_base_id: str, *, limit: int = 200) -> dict[str, Any]:
        value = _data(
            await self._request(
                "GET",
                f"/api/v1/knowledgebase/{quote(knowledge_base_id, safe='')}/wiki/graph",
                params={"mode": "overview", "limit": max(1, min(limit, 10000))},
            )
        )
        if not isinstance(value, dict):
            raise WeKnoraError("WeKnora returned an invalid wiki graph")
        return dict(value)

    async def get_wiki_stats(self, knowledge_base_id: str) -> dict[str, Any]:
        value = _data(await self._request("GET", f"/api/v1/knowledgebase/{quote(knowledge_base_id, safe='')}/wiki/stats"))
        if not isinstance(value, dict):
            raise WeKnoraError("Invalid Wiki generation status")
        return value

    async def create_wiki_folder(self, knowledge_base_id: str, *, name: str, parent_id: str = "") -> dict[str, Any]:
        value = _data(await self._request("POST", f"/api/v1/knowledgebase/{quote(knowledge_base_id, safe='')}/wiki/folders", json={"name": name, "parent_id": parent_id}))
        if not isinstance(value, dict) or not value.get("id"):
            raise WeKnoraError("Invalid Wiki folder response")
        return value

    async def rebuild_wiki_links(self, knowledge_base_id: str) -> None:
        await self._request("POST", f"/api/v1/knowledgebase/{quote(knowledge_base_id, safe='')}/wiki/rebuild-links")

    async def list_wiki_pages(
        self,
        knowledge_base_id: str,
        *,
        page: int = 1,
        page_size: int = 50,
        query: str = "",
    ) -> dict[str, Any]:
        payload = await self._request(
            "GET",
            f"/api/v1/knowledgebase/{quote(knowledge_base_id, safe='')}/wiki/pages",
            params={"page": page, "page_size": page_size, "query": query or None},
        )
        value = _data(payload)
        source = value if isinstance(value, dict) else payload if isinstance(payload, dict) else {}
        rows = _as_list(source.get("pages") if isinstance(source, dict) else value)
        return {
            "pages": rows,
            "total": int(source.get("total") or len(rows)),
            "page": int(source.get("page") or page),
            "page_size": int(source.get("page_size") or page_size),
            "total_pages": int(source.get("total_pages") or max(1, (int(source.get("total") or len(rows)) + page_size - 1) // page_size)),
        }

    async def get_wiki_page(self, knowledge_base_id: str, slug: str) -> dict[str, Any]:
        safe_slug = "/".join(quote(part, safe="") for part in slug.split("/"))
        value = _data(
            await self._request(
                "GET",
                f"/api/v1/knowledgebase/{quote(knowledge_base_id, safe='')}/wiki/pages/{safe_slug}",
            )
        )
        if not isinstance(value, dict):
            raise WeKnoraError("WeKnora returned an invalid Wiki page")
        return dict(value)

    async def get_wiki_index(
        self,
        knowledge_base_id: str,
        *,
        types: Sequence[str],
        limit: int = 50,
        cursor: str = "",
    ) -> dict[str, Any]:
        params: dict[str, Any] = {
            "types": ",".join(value for value in types if value),
            "limit": max(1, min(limit, 500)),
        }
        if cursor:
            params["cursor"] = cursor
        value = _data(
            await self._request(
                "GET",
                f"/api/v1/knowledgebase/{quote(knowledge_base_id, safe='')}/wiki/index",
                params=params,
            )
        )
        if not isinstance(value, dict):
            raise WeKnoraError("WeKnora returned an invalid Wiki index")
        return dict(value)

    async def create_wiki_page(self, knowledge_base_id: str, payload: Mapping[str, Any]) -> dict[str, Any]:
        value = _data(
            await self._request(
                "POST",
                f"/api/v1/knowledgebase/{quote(knowledge_base_id, safe='')}/wiki/pages",
                json=dict(payload),
            )
        )
        if not isinstance(value, dict):
            raise WeKnoraError("WeKnora returned an invalid Wiki page")
        return dict(value)

    async def update_wiki_page(self, knowledge_base_id: str, slug: str, changes: Mapping[str, Any]) -> dict[str, Any]:
        safe_slug = "/".join(quote(part, safe="") for part in slug.split("/"))
        value = _data(
            await self._request(
                "PUT",
                f"/api/v1/knowledgebase/{quote(knowledge_base_id, safe='')}/wiki/pages/{safe_slug}",
                json=dict(changes),
            )
        )
        return dict(value) if isinstance(value, dict) else {}

    async def delete_wiki_page(self, knowledge_base_id: str, slug: str) -> None:
        safe_slug = "/".join(quote(part, safe="") for part in slug.split("/"))
        await self._request(
            "DELETE",
            f"/api/v1/knowledgebase/{quote(knowledge_base_id, safe='')}/wiki/pages/{safe_slug}",
        )

    async def search(self, query: str, knowledge_base_ids: Sequence[str]) -> list[dict[str, Any]]:
        ids = [str(value) for value in knowledge_base_ids if str(value)]
        if not ids:
            return []
        rows = _as_list(
            _data(
                await self._request(
                    "POST",
                    "/api/v1/knowledge-search",
                    json={"query": query, "knowledge_base_ids": ids},
                )
            )
        )
        normalized: list[dict[str, Any]] = []
        for row in rows:
            normalized.append(
                {
                    "chunk_id": str(row.get("id") or row.get("chunk_id") or ""),
                    "content": str(row.get("content") or ""),
                    "knowledge_id": str(row.get("knowledge_id") or ""),
                    "knowledge_base_id": str(row.get("knowledge_base_id") or ""),
                    "title": str(row.get("knowledge_title") or row.get("title") or row.get("knowledge_filename") or ""),
                    "filename": str(row.get("knowledge_filename") or ""),
                    "score": float(row.get("score") or 0),
                    "chunk_index": int(row.get("chunk_index") or 0),
                    "metadata": row.get("metadata") if isinstance(row.get("metadata"), dict) else {},
                }
            )
        return normalized

    async def get_chunk(self, chunk_id: str) -> dict[str, Any]:
        value = _data(await self._request("GET", f"/api/v1/chunks/by-id/{quote(chunk_id, safe='')}"))
        if not isinstance(value, dict):
            raise WeKnoraError("WeKnora returned an invalid chunk")
        return dict(value)

    async def get_chunk_context(self, chunk_id: str, *, expected_knowledge_base_id: str, radius: int = 2) -> list[dict[str, Any]]:
        center = await self.get_chunk(chunk_id)
        if str(center.get("knowledge_base_id") or "") != expected_knowledge_base_id:
            raise WeKnoraError("Chunk does not belong to the authorized knowledge base", status_code=404)

        before: list[dict[str, Any]] = []
        cursor = center
        for _ in range(max(0, min(radius, 4))):
            previous_id = str(cursor.get("pre_chunk_id") or "")
            if not previous_id:
                break
            previous = await self.get_chunk(previous_id)
            if str(previous.get("knowledge_base_id") or "") != expected_knowledge_base_id:
                break
            before.append(previous)
            cursor = previous

        after: list[dict[str, Any]] = []
        cursor = center
        for _ in range(max(0, min(radius, 4))):
            next_id = str(cursor.get("next_chunk_id") or "")
            if not next_id:
                break
            following = await self.get_chunk(next_id)
            if str(following.get("knowledge_base_id") or "") != expected_knowledge_base_id:
                break
            after.append(following)
            cursor = following
        return [*reversed(before), center, *after]