"""MemoryClient -- Python SDK for the 9D Labs Memory Runtime v1 API.

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

Usage::

    from nined.memory import MemoryClient

    client = MemoryClient(base_url="https://api.9dlabs.xyz", api_key="your-key")

    # Ingest
    client.ingest(
        workspace_id="ws-1",
        actor_id="agent-1",
        artifact_type="document",
        raw_payload={"title": "Runbook", "content": "If auth fails, restart the pod."},
    )

    # Retrieve context
    pack = client.context_pack(query="What should I do about auth?", workspace_id="ws-1")
    print(pack["context_text"])

    # Audit trail
    receipt = client.get_receipt(pack["receipt_id"])
"""

from __future__ import annotations

from typing import Any

from nined.memory._transport import BaseTransport


class MemoryClient(BaseTransport):
    """Client for the Memory Runtime v1 API.

    Args:
        base_url: Server URL (e.g. ``https://api.9dlabs.xyz``).
        api_key:  API key for authentication (X-API-Key header).
        timeout:  HTTP timeout in seconds (default 30).
    """

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

    # -- Ingest ---------------------------------------------------------------

    def ingest(
        self,
        *,
        workspace_id: str,
        actor_id: str,
        artifact_type: str,
        raw_payload: dict[str, Any] | str | list[Any] | None = None,
        blob_pointer: str | None = None,
        permissions: dict[str, Any] | None = None,
        idempotency_key: str | None = None,
        session_id: str | None = None,
        thread_id: str | None = None,
        timestamp: str | None = None,
        callback_url: str | None = None,
    ) -> dict[str, Any]:
        """Ingest a single artifact into workspace memory.

        Returns:
            ``{"artifact_id": ..., "content_hash": ..., "indexed": bool}``
        """
        payload: dict[str, Any] = {
            "workspace_id": workspace_id,
            "actor_id": actor_id,
            "artifact_type": artifact_type,
        }
        if raw_payload is not None:
            payload["raw_payload"] = raw_payload
        if blob_pointer is not None:
            payload["blob_pointer"] = blob_pointer
        if permissions is not None:
            payload["permissions"] = permissions
        if idempotency_key is not None:
            payload["idempotency_key"] = idempotency_key
        if session_id is not None:
            payload["session_id"] = session_id
        if thread_id is not None:
            payload["thread_id"] = thread_id
        if timestamp is not None:
            payload["timestamp"] = timestamp
        if callback_url is not None:
            payload["callback_url"] = callback_url
        return self._post("/v1/ingest", payload)

    # -- Context Pack ---------------------------------------------------------

    def context_pack(
        self,
        *,
        query: str,
        workspace_id: str,
        max_tokens: int = 4096,
        actor_id: str | None = None,
        mode: str | None = None,
        max_latency_ms: int | None = None,
        max_cost_usd: float | None = None,
        min_evidence_tokens: int | None = None,
        min_confidence: float | None = None,
        max_spans_per_artifact: int | None = None,
        diversity_decay: float | None = None,
        redundancy_jaccard_threshold: float | None = None,
        redundancy_heavy_overlap_threshold: float | None = None,
        temporal_half_life_hours: float | None = None,
    ) -> dict[str, Any]:
        """Retrieve a deterministic context pack for a query.

        Returns:
            ``{"context_text": ..., "citations": [...], "confidence": float,
            "receipt_id": ..., "total_tokens": int, "pack_hash": ...,
            "abstain_flag": bool, ...}``
        """
        payload: dict[str, Any] = {
            "query": query,
            "workspace_id": workspace_id,
            "max_tokens": max_tokens,
        }
        _set_optional(payload, "actor_id", actor_id)
        _set_optional(payload, "mode", mode)
        _set_optional(payload, "max_latency_ms", max_latency_ms)
        _set_optional(payload, "max_cost_usd", max_cost_usd)
        _set_optional(payload, "min_evidence_tokens", min_evidence_tokens)
        _set_optional(payload, "min_confidence", min_confidence)
        _set_optional(payload, "max_spans_per_artifact", max_spans_per_artifact)
        _set_optional(payload, "diversity_decay", diversity_decay)
        _set_optional(payload, "redundancy_jaccard_threshold", redundancy_jaccard_threshold)
        _set_optional(payload, "redundancy_heavy_overlap_threshold", redundancy_heavy_overlap_threshold)
        _set_optional(payload, "temporal_half_life_hours", temporal_half_life_hours)
        return self._post("/v1/context-pack", payload)

    # -- Feedback -------------------------------------------------------------

    def feedback(
        self,
        *,
        workspace_id: str,
        actor_id: str,
        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 (pin, private, mark_wrong, outdated, resolve_conflict).

        Returns:
            ``{"feedback_id": ..., "action": ..., "applied": True}``
        """
        payload: dict[str, Any] = {
            "workspace_id": workspace_id,
            "actor_id": actor_id,
            "action": action,
        }
        _set_optional(payload, "artifact_id", artifact_id)
        _set_optional(payload, "span_id", span_id)
        _set_optional(payload, "conflict_id", conflict_id)
        _set_optional(payload, "preferred_span_id", preferred_span_id)
        _set_optional(payload, "reason", reason)
        return self._post("/v1/feedback", payload)

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

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

        Returns:
            ``{"receipts": [...], "total_count": int, "limit": int, "offset": int}``
        """
        params: dict[str, str] = {"limit": str(limit), "offset": str(offset)}
        if workspace_id is not None:
            params["workspace_id"] = workspace_id
        return self._get("/v1/receipts", params=params)

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

        Returns:
            ``{"id": ..., "query": ..., "receipt": {...}, ...}``
        """
        return self._get(f"/v1/receipts/{receipt_id}")

    # -- 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, "total_tokens": int, ...}``
        """
        return self._get(f"/v1/artifacts/{artifact_id}/status")

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

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

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

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

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

        Returns:
            ``{"answer": ..., "confidence": float, "citations": [...],
            "receipt_id": ..., "abstain_flag": bool, ...}``
        """
        return self._post("/v1/study/ask", {
            "query": query,
            "workspace_id": workspace_id,
            "max_tokens": max_tokens,
        })

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

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

    # -- Registration & Billing -----------------------------------------------

    def register(self, email: str) -> dict[str, Any]:
        """Register a new tenant (multi-tenant mode only).

        Returns:
            ``{"tenant_id": ..., "api_key": ..., "message": ...}``
        """
        return self._post("/v1/register", {"email": email})

    def regenerate_key(self, email: str) -> dict[str, Any]:
        """Regenerate API key for an existing tenant.

        Returns:
            ``{"tenant_id": ..., "api_key": ..., "message": ...}``
        """
        return self._post("/v1/regenerate-key", {"email": email})

    def billing_usage(self) -> dict[str, Any]:
        """Get current billing usage for the authenticated tenant."""
        return self._get("/v1/billing/usage")

    def billing_plan(self) -> dict[str, Any]:
        """Get current plan and entitlements for the authenticated tenant."""
        return self._get("/v1/billing/plan")


def _set_optional(payload: dict[str, Any], key: str, value: Any) -> None:
    if value is not None:
        payload[key] = value
