"""
Agent Storage — Persistent shared state for agents.

Usage from within an agent:

    from orchagent import storage

    # Write a document
    storage.set("signals", "2026-03-05", {"pending": [...], "done": [...]})

    # Read a document
    data = storage.get("signals", "2026-03-05")

    # Atomic read-modify-write (with CAS retry)
    storage.update("signals", "2026-03-05", lambda doc: {**doc, "done": doc["done"] + [new]})

    # List keys in a namespace
    result = storage.list("signals")  # {"keys": [...], "cursor": ..., "has_more": ...}

    # List all namespaces
    namespaces = storage.namespaces()

    # Delete
    storage.delete("signals", "2026-03-05")
    storage.delete_namespace("signals")

All functions are synchronous (blocking) for simplicity in agent scripts.
They use httpx under the hood to call the gateway REST API.
"""

from __future__ import annotations

import os
import time
from typing import Any, Callable, Optional

import httpx

# Max retries for CAS (compare-and-swap) conflicts
_CAS_MAX_RETRIES = 5
_CAS_BASE_DELAY = 0.1  # seconds

_DEFAULT_GATEWAY_URL = "https://api.orchagent.io"
_DEFAULT_TIMEOUT = 30.0


class StorageError(Exception):
    """Base exception for storage operations."""
    def __init__(self, message: str, status_code: int | None = None):
        super().__init__(message)
        self.status_code = status_code


class StorageNotFoundError(StorageError):
    """Document or namespace not found."""
    pass


class StorageConflictError(StorageError):
    """CAS version conflict (concurrent modification)."""
    pass


class StorageQuotaError(StorageError):
    """Storage quota exceeded."""
    pass


def _gateway_url() -> str:
    return os.environ.get("ORCHAGENT_GATEWAY_URL", _DEFAULT_GATEWAY_URL).rstrip("/")


def _auth_headers() -> dict[str, str]:
    key = os.environ.get("ORCHAGENT_SERVICE_KEY") or os.environ.get("ORCHAGENT_API_KEY", "")
    headers: dict[str, str] = {}
    if key:
        headers["Authorization"] = f"Bearer {key}"
    return headers


def _request(
    method: str,
    path: str,
    *,
    json_body: Any = None,
    headers: dict[str, str] | None = None,
) -> httpx.Response:
    """Make an HTTP request to the gateway storage API."""
    url = f"{_gateway_url()}/storage{path}"
    req_headers = {**_auth_headers(), **(headers or {})}

    with httpx.Client(timeout=_DEFAULT_TIMEOUT) as client:
        response = client.request(
            method,
            url,
            json=json_body,
            headers=req_headers,
        )

    return response


def _handle_error(response: httpx.Response) -> None:
    """Raise appropriate exception for error responses."""
    if response.status_code < 400:
        return

    try:
        data = response.json()
        message = data.get("error", {}).get("message", response.text)
        code = data.get("error", {}).get("code", "")
    except Exception:
        message = response.text
        code = ""

    if response.status_code == 404:
        raise StorageNotFoundError(message, response.status_code)
    elif response.status_code == 409:
        raise StorageConflictError(message, response.status_code)
    elif response.status_code == 413:
        raise StorageQuotaError(message, response.status_code)
    elif response.status_code == 429 and ("QUOTA" in code or code.startswith("STORAGE_")):
        raise StorageQuotaError(message, response.status_code)
    else:
        raise StorageError(message, response.status_code)


def _get_document(namespace: str, key: str) -> Optional[dict]:
    """
    Read the full document envelope from the API.

    Returns None when document does not exist.
    """
    response = _request("GET", f"/{namespace}/{key}")
    if response.status_code == 404:
        return None
    _handle_error(response)
    return response.json()


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def get(namespace: str, key: str) -> Optional[Any]:
    """
    Read a document. Returns None if not found.

    Returns only the stored JSON value.
    """
    doc = _get_document(namespace, key)
    if doc is None:
        return None
    return doc.get("value")


def set(namespace: str, key: str, value: Any, *, version: int | None = None) -> dict:
    """
    Write a document (create or update).

    Args:
        namespace: Namespace string (e.g. "signals")
        key: Document key (e.g. "2026-03-05")
        value: JSON-serializable value (dict or list)
        version: If set, performs CAS — fails with StorageConflictError if
                 the document was modified since this version.

    Returns the written document metadata.
    """
    headers: dict[str, str] = {}
    if version is not None:
        headers["If-Match"] = str(version)

    response = _request("PUT", f"/{namespace}/{key}", json_body=value, headers=headers)
    _handle_error(response)
    return response.json()


def update(
    namespace: str,
    key: str,
    fn: Callable[[Optional[Any]], Any],
    *,
    max_retries: int = _CAS_MAX_RETRIES,
) -> dict:
    """
    Atomic read-modify-write with CAS retry.

    fn receives the current document value (or None if it doesn't exist)
    and must return the new value to write.

    Automatically retries on version conflicts up to max_retries times.
    """
    for attempt in range(max_retries):
        doc = _get_document(namespace, key)
        current_value = doc.get("value") if doc else None
        current_version = doc["version"] if doc else None

        new_value = fn(current_value)

        try:
            return set(namespace, key, new_value, version=current_version)
        except StorageConflictError:
            if attempt == max_retries - 1:
                raise
            delay = _CAS_BASE_DELAY * (2 ** attempt)
            time.sleep(delay)

    raise StorageConflictError("Max CAS retries exceeded")


def patch(namespace: str, key: str, partial: dict) -> dict:
    """
    Atomic merge-patch: shallow-merge partial into existing document.

    Raises StorageNotFoundError if document doesn't exist.
    """
    response = _request("PATCH", f"/{namespace}/{key}", json_body=partial)
    _handle_error(response)
    return response.json()


def list(namespace: str, *, limit: int = 100, cursor: str | None = None) -> dict:
    """
    List keys in a namespace.

    Returns: {"keys": [...], "cursor": str|None, "has_more": bool}
    """
    params = f"?limit={limit}"
    if cursor:
        params += f"&cursor={cursor}"
    response = _request("GET", f"/{namespace}{params}")
    _handle_error(response)
    return response.json()


def namespaces() -> list[str]:
    """List all namespaces in the workspace."""
    response = _request("GET", "")
    _handle_error(response)
    return response.json().get("namespaces", [])


def delete(namespace: str, key: str) -> bool:
    """Delete a document. Returns True if deleted, raises StorageNotFoundError if not found."""
    response = _request("DELETE", f"/{namespace}/{key}")
    _handle_error(response)
    return True


def delete_namespace(namespace: str) -> int:
    """Delete all documents in a namespace. Returns count deleted."""
    response = _request("DELETE", f"/{namespace}")
    _handle_error(response)
    return response.json().get("documents_deleted", 0)
