"""MemoryClientV2 -- Python SDK for the 9D Labs Memory Runtime v2 API.

Zero external dependencies. Uses only stdlib ``urllib`` and ``json``.

Usage::

    from nined.memory import MemoryClientV2

    client = MemoryClientV2(
        base_url="https://api.9dlabs.xyz",
        api_key="your-key",
        workspace_id="my-workspace",
    )

    # Batch ingest
    client.ingest([
        {"artifact_type": "document", "raw_payload": {"content": "..."}},
        {"artifact_type": "note", "raw_payload": {"content": "..."}},
    ])

    # Retrieve context with serving profile
    pack = client.context_pack("What happened?", profile="high_recall")
    for snippet in pack["snippets"]:
        print(snippet["content"])
"""

from __future__ import annotations

from typing import Any

from nined.memory._transport import BaseTransport


class MemoryClientV2(BaseTransport):
    """Client for the Memory Runtime v2 API.

    Args:
        base_url:     Server URL (e.g. ``https://api.9dlabs.xyz``).
        api_key:      API key for authentication.
        workspace_id: Default workspace scope for all operations.
        actor_id:     Default actor identity for all operations.
        timeout:      HTTP timeout in seconds (default 30).
    """

    def __init__(
        self,
        *,
        base_url: str = "https://api.9dlabs.xyz",
        api_key: str = "",
        workspace_id: str,
        actor_id: str = "sdk",
        timeout: int = 30,
    ):
        self.base_url = base_url.rstrip("/")
        self.api_key = api_key
        self.workspace_id = workspace_id
        self.actor_id = actor_id
        self.timeout = timeout

    def _v2_headers(self, *, auth: bool = True) -> dict[str, str]:
        headers = self._headers(auth=auth)
        if auth:
            headers["x-api-key"] = headers.pop("X-API-Key", "")
            headers["x-workspace-id"] = self.workspace_id
            headers["x-actor-id"] = self.actor_id
        return headers

    def _v2_post(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
        import json
        url = f"{self.base_url}{path}"
        data = json.dumps(payload).encode("utf-8")
        from nined.memory._transport import send_request
        return send_request(
            method="POST", url=url, headers=self._v2_headers(),
            data=data, timeout=self.timeout,
        )

    def _v2_get(
        self, path: str, *, params: dict[str, str] | None = None, auth: bool = True,
    ) -> dict[str, Any]:
        import urllib.parse
        url = f"{self.base_url}{path}"
        if params:
            url = f"{url}?{urllib.parse.urlencode(params)}"
        from nined.memory._transport import send_request
        return send_request(
            method="GET", url=url, headers=self._v2_headers(auth=auth), timeout=self.timeout,
        )

    def _v2_delete(self, path: str) -> dict[str, Any]:
        from nined.memory._transport import send_request
        return send_request(
            method="DELETE", url=f"{self.base_url}{path}",
            headers=self._v2_headers(), timeout=self.timeout,
        )

    # -- Memory ---------------------------------------------------------------

    def ingest(
        self,
        artifacts: list[dict[str, Any]],
        *,
        async_index: bool = False,
    ) -> dict[str, Any]:
        """Ingest one or more artifacts into workspace memory.

        Args:
            artifacts:   List of artifact dicts (``artifact_type``, ``raw_payload``, etc).
            async_index: If True, artifacts are queued for background indexing.

        Returns:
            Sync: ``{"workspace_id": ..., "ingested": [{artifact_id, span_ids, created}]}``
            Async: ``{"workspace_id": ..., "queued_jobs": [{job_id, queued}]}``
        """
        return self._v2_post("/v2/ingest", {
            "workspace_id": self.workspace_id,
            "actor_id": self.actor_id,
            "async_index": async_index,
            "artifacts": artifacts,
        })

    def context_pack(
        self,
        query: str,
        *,
        max_tokens: int = 4096,
        profile: str | None = None,
    ) -> dict[str, Any]:
        """Retrieve a deterministic context pack for a query.

        Args:
            query:      Natural-language question.
            max_tokens: Token budget for the pack.
            profile:    Serving profile (``low_latency``, ``balanced``, ``high_recall``).

        Returns:
            ``{"snippets": [...], "pack_hash": ..., "token_accounting": {...},
            "serving_profile": ..., "citation_count": int, ...}``
        """
        payload: dict[str, Any] = {
            "workspace_id": self.workspace_id,
            "actor_id": self.actor_id,
            "query": query,
            "max_tokens": max_tokens,
        }
        if profile is not None:
            payload["profile"] = profile
        return self._v2_post("/v2/context-pack", payload)

    def feedback(
        self,
        action: str,
        *,
        artifact_id: str | None = None,
        span_id: str | None = None,
        conflict_id: str | None = None,
        preferred_span_id: str | None = None,
        reason: str | None = None,
    ) -> dict[str, Any]:
        """Submit a correction event.

        Returns:
            ``{"feedback_id": ..., "applied": True, "action": ...}``
        """
        payload: dict[str, Any] = {
            "workspace_id": self.workspace_id,
            "actor_id": self.actor_id,
            "action": action,
        }
        if artifact_id is not None:
            payload["artifact_id"] = artifact_id
        if span_id is not None:
            payload["span_id"] = span_id
        if conflict_id is not None:
            payload["conflict_id"] = conflict_id
        if preferred_span_id is not None:
            payload["preferred_span_id"] = preferred_span_id
        if reason is not None:
            payload["reason"] = reason
        return self._v2_post("/v2/feedback", payload)

    # -- Receipts -------------------------------------------------------------

    def list_receipts(self, *, limit: int = 50, offset: int = 0) -> dict[str, Any]:
        """List decision receipts for the workspace.

        Returns:
            ``{"workspace_id": ..., "count": int, "receipts": [...]}``
        """
        params = {"limit": str(limit), "offset": str(offset)}
        return self._v2_get(f"/v2/receipts/{self.workspace_id}", params=params)

    def get_receipt(self, pack_hash: str) -> dict[str, Any]:
        """Get a single receipt by pack hash.

        Returns:
            Full receipt with selection metadata and diagnostics.
        """
        return self._v2_get(f"/v2/receipts/{self.workspace_id}/{pack_hash}")

    # -- Status ---------------------------------------------------------------

    def artifact_status(self, artifact_id: str) -> dict[str, Any]:
        """Get indexing status for an artifact.

        Returns:
            ``{"artifact_id": ..., "indexed": bool, "span_count": int, ...}``
        """
        params = {"workspace_id": self.workspace_id}
        return self._v2_get(f"/v2/artifacts/{artifact_id}/status", params=params)

    def job_status(self, job_id: str) -> dict[str, Any]:
        """Get status of an async indexing job.

        Returns:
            ``{"job_id": ..., "state": ..., "attempts": int, ...}``
        """
        return self._v2_get(f"/v2/jobs/{job_id}")

    # -- Study ----------------------------------------------------------------

    def ask(
        self,
        query: str,
        *,
        max_tokens: int = 16000,
        profile: str | None = None,
    ) -> dict[str, Any]:
        """Get an LLM-synthesized answer from workspace memory.

        Returns:
            ``{"answer": ..., "confidence": float, "citations": [...],
            "pack_hash": ..., "abstain_flag": bool, ...}``
        """
        payload: dict[str, Any] = {
            "workspace_id": self.workspace_id,
            "actor_id": self.actor_id,
            "query": query,
            "max_tokens": max_tokens,
        }
        if profile is not None:
            payload["profile"] = profile
        return self._v2_post("/v2/study/ask", payload)

    # -- Workspace ------------------------------------------------------------

    def delete_workspace(self) -> dict[str, Any]:
        """Delete all data in the workspace.

        Returns:
            ``{"workspace_id": ..., "deleted_artifacts": int}``
        """
        return self._v2_delete(f"/v2/workspace/{self.workspace_id}")

    # -- Health ---------------------------------------------------------------

    def health(self) -> dict[str, Any]:
        """Check service liveness."""
        return self._v2_get("/v2/health", auth=False)

    def ready(self) -> dict[str, Any]:
        """Check service readiness including storage health."""
        return self._v2_get("/v2/ready", auth=False)
