# SPDX-FileCopyrightText: 2026 Surendra Singh <66422685+bolnet@users.noreply.github.com>
# SPDX-License-Identifier: MIT
"""Starlette ASGI REST API for the Attestor memory service."""

from __future__ import annotations

import logging
import os
import threading
from typing import Any

from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import JSONResponse
from starlette.routing import Route

logger = logging.getLogger("attestor.api")

# Lazy singleton — initialized on first request. The lock prevents a
# TOCTOU race under threaded ASGI servers (Uvicorn workers, mixed
# sync/async handler pools) where two concurrent first-callers both
# observe ``_mem is None``, both build an AgentMemory, and the second
# overwrites the first — leaking open Postgres connections.
_mem = None
_mem_lock = threading.Lock()


def _pinecone_config_from_env() -> dict[str, Any] | None:
    """Resolve the Pinecone (vector-role) backend config for the env path.

    The env path (POSTGRES_URL / NEO4J_URI) is for cloud deploys that bind
    secrets via env. It used to register only postgres (document) + neo4j
    (graph) and silently drop the vector role — so semantic recall degraded
    to tag+graph on every env-driven deploy. (The vector role moved from the
    Postgres pgvector column to a dedicated Pinecone backend on 2026-04-29;
    this path was never updated, so ``recall`` ran with ``_vector_store=None``.)

    The YAML stack is the source of truth for the non-secret index settings
    (index_name / cloud / region / metric / dimension). The secret (api_key)
    and connection host may be overridden via ``PINECONE_API_KEY`` /
    ``PINECONE_HOST`` / ``PINECONE_INDEX`` so cloud deploys never bake the key
    into a config file.

    Returns the pinecone config dict, or None when neither YAML nor env
    defines a vector backend — legacy pgvector-only deploys then keep running
    without a separate vector store (the prior behavior).
    """
    pcn: dict[str, Any] = {}

    try:
        from attestor.config import get_stack

        stack = get_stack()
        if stack.pinecone is not None:
            pcn = {
                "index_name": stack.pinecone.index_name,
                "metric": stack.pinecone.metric,
                "cloud": stack.pinecone.cloud,
                "region": stack.pinecone.region,
                "dimension": stack.embedder.dimensions,
            }
            if stack.pinecone.host:
                pcn["host"] = stack.pinecone.host
            yaml_key = os.environ.get(stack.pinecone.api_key_env)
            if yaml_key:
                pcn["api_key"] = yaml_key
    except (Exception, SystemExit):  # noqa: BLE001 — YAML may be absent or
        # missing required env on a pure-env deploy; the loader raises
        # SystemExit (not Exception) in that case. Fall back to PINECONE_*.
        pass

    # Env overrides win over YAML for the secret + connection target.
    api_key = os.environ.get("PINECONE_API_KEY")
    host = os.environ.get("PINECONE_HOST")
    index = os.environ.get("PINECONE_INDEX")
    if api_key:
        pcn["api_key"] = api_key
    if host:
        pcn["host"] = host
    if index:
        pcn["index_name"] = index

    # Register the vector backend only when we can actually reach an index.
    # An index_name is the minimum; without it there is nothing to query.
    if not pcn.get("index_name"):
        return None
    return pcn


def _build_config() -> dict[str, Any] | None:
    """Build backend config. Returns None when no backend can be resolved.

    Resolution order (top wins):
        1. Explicit env vars: POSTGRES_URL [+ NEO4J_URI] / NEO4J_URI alone —
           preserved for cloud deploys that bind secrets via env. The vector
           role (Pinecone) is folded in from the YAML SoT + PINECONE_* env so
           it is never silently dropped.
        2. ``configs/attestor.yaml`` via ``attestor.config.get_stack()`` —
           when no env vars are set we honor the YAML SoT. Pass
           ATTESTOR_CONFIG=/path/to/yaml to override the default file.
        3. None — caller bubbles a clear error.
    """
    postgres_url = os.environ.get("POSTGRES_URL")
    neo4j_uri = os.environ.get("NEO4J_URI")

    cfg: dict[str, Any] = {}
    backends: list = []

    if postgres_url:
        backends.append("postgres")
        cfg["postgres"] = {
            "mode": "cloud",
            "url": postgres_url,
            "database": os.environ.get("POSTGRES_DATABASE", "attestor"),
            "auth": {
                "username": os.environ.get("POSTGRES_USERNAME", "postgres"),
                "password": os.environ.get("POSTGRES_PASSWORD", ""),
            },
            "sslmode": os.environ.get("POSTGRES_SSLMODE"),
        }

    if neo4j_uri:
        backends.append("neo4j")
        cfg["neo4j"] = {
            "mode": "cloud",
            "url": neo4j_uri,
            "database": os.environ.get("NEO4J_DATABASE", "neo4j"),
            "auth": {
                "username": os.environ.get("NEO4J_USERNAME", "neo4j"),
                "password": os.environ.get("NEO4J_PASSWORD", ""),
            },
        }

    if backends:
        # On the env-driven path, preserve the Pinecone vector role. Without
        # this the vector store is silently absent and recall loses its
        # semantic lane (see _pinecone_config_from_env for the regression).
        pinecone_cfg = _pinecone_config_from_env()
        if pinecone_cfg is not None:
            backends.append("pinecone")
            cfg["pinecone"] = pinecone_cfg
        cfg["backends"] = backends
        return cfg

    # No env override — defer to the YAML SoT.
    try:
        from attestor.config import build_backend_config, get_stack
        return build_backend_config(get_stack())
    except SystemExit:
        # YAML missing required env (e.g. NEO4J_PASSWORD) — return None so
        # the api surfaces a clear "configure backends" error rather than
        # the loader's exit message.
        return None


def _get_mem():
    global _mem
    if _mem is not None:
        return _mem
    with _mem_lock:
        # Double-checked locking: re-test inside the lock so we don't
        # rebuild after another thread populated _mem while we waited.
        if _mem is None:
            from attestor.core import AgentMemory
            from attestor._paths import resolve_data_dir

            data_dir = resolve_data_dir()
            config = _build_config()
            if config is not None:
                _mem = AgentMemory(data_dir, config=config)
            else:
                _mem = AgentMemory(data_dir)
    return _mem


def _ok(data: Any) -> JSONResponse:
    return JSONResponse({"ok": True, "data": data})


def _err(msg: str, status: int = 400) -> JSONResponse:
    return JSONResponse({"ok": False, "error": msg}, status_code=status)


async def health(request: Request) -> JSONResponse:
    mem = _get_mem()
    return _ok(mem.health())


async def stats(request: Request) -> JSONResponse:
    mem = _get_mem()
    return _ok(mem.stats())


async def add_memory(request: Request) -> JSONResponse:
    body = await request.json()
    content = body.get("content")
    if not content:
        return _err("content is required")
    mem = _get_mem()
    # Import lazily — quotas is Postgres-specific and may pull psycopg2.
    from attestor.quotas import QuotaExceeded

    try:
        m = mem.add(
            content=content,
            tags=body.get("tags", []),
            category=body.get("category", "general"),
            entity=body.get("entity"),
            namespace=body.get("namespace", "default"),
            event_date=body.get("event_date"),
            confidence=body.get("confidence", 1.0),
            metadata=body.get("metadata", {}),
        )
    except QuotaExceeded as e:
        # 429 Too Many Requests — include the limit field name so clients
        # can show a useful error (e.g. "max_writes_per_day reached").
        return _err(f"quota exceeded: {e}", status=429)
    return _ok(m.to_dict())


async def recall(request: Request) -> JSONResponse:
    body = await request.json()
    query = body.get("query")
    if not query:
        return _err("query is required")
    mem = _get_mem()
    results = mem.recall(
        query, budget=body.get("budget"), namespace=body.get("namespace")
    )
    return _ok([{"content": r.memory.content, "score": r.score,
                 "source": r.match_source, "id": r.memory.id,
                 "memory": r.memory.to_dict()} for r in results])


async def search(request: Request) -> JSONResponse:
    body = await request.json()
    mem = _get_mem()
    memories = mem.search(
        query=body.get("query"),
        category=body.get("category"),
        entity=body.get("entity"),
        namespace=body.get("namespace"),
        status=body.get("status", "active"),
        after=body.get("after"),
        before=body.get("before"),
        limit=body.get("limit", 10),
    )
    return _ok([m.to_dict() for m in memories])


async def timeline(request: Request) -> JSONResponse:
    body = await request.json()
    entity = body.get("entity")
    if not entity:
        return _err("entity is required")
    mem = _get_mem()
    memories = mem.timeline(entity, namespace=body.get("namespace"))
    return _ok([m.to_dict() for m in memories])


async def forget(request: Request) -> JSONResponse:
    body = await request.json()
    memory_id = body.get("memory_id")
    if not memory_id:
        return _err("memory_id is required")
    mem = _get_mem()

    # AuthZ: when the JWTAuthMiddleware is mounted (HOSTED / SHARED modes)
    # ``request.state.user_id`` is populated. Reject forget calls that
    # target a memory the caller doesn't own — otherwise any authenticated
    # user could enumerate memory_ids and archive other users' data.
    caller_user_id = getattr(request.state, "user_id", None)
    if caller_user_id is not None:
        target = mem.get(memory_id)
        if target is None:
            return _err("not found", status=404)
        owner = getattr(target, "user_id", None)
        if owner is not None and owner != caller_user_id:
            return _err("forbidden", status=403)

    ok = mem.forget(memory_id)
    return _ok({"forgotten": ok})


async def get_memory(request: Request) -> JSONResponse:
    memory_id = request.path_params["memory_id"]
    mem = _get_mem()
    m = mem.get(memory_id)
    if m is None:
        return _err("not found", status=404)
    return _ok(m.to_dict())


routes = [
    Route("/health", health, methods=["GET"]),
    Route("/stats", stats, methods=["GET"]),
    Route("/add", add_memory, methods=["POST"]),
    Route("/recall", recall, methods=["POST"]),
    Route("/search", search, methods=["POST"]),
    Route("/timeline", timeline, methods=["POST"]),
    Route("/forget", forget, methods=["POST"]),
    Route("/memory/{memory_id}", get_memory, methods=["GET"]),
]

# Attach the read-only UI at /ui/* (routes have absolute paths)
try:
    from attestor.ui.app import ui_routes

    routes.extend(ui_routes())
except ImportError as e:  # pragma: no cover - UI is optional
    logger.warning("UI router unavailable: %s", e)

# Attach the audit explorer at /audit/* + the static /audit.html dashboard.
# Off by default — opt-in via `ATTESTOR_AUDIT_ENABLED=1` or
# `audit.dashboard.enabled: true` in configs/attestor.yaml. Mounting only
# when enabled keeps the surface smaller in single-user installs and lets
# operators gate the audit trail behind explicit configuration.
try:
    from attestor.audit.routes import audit_routes, is_enabled as _audit_enabled

    if _audit_enabled():
        routes.extend(audit_routes())
except ImportError as e:  # pragma: no cover - audit module is optional
    logger.warning("Audit router unavailable: %s", e)


def _build_middleware() -> list:
    """Return the Starlette middleware stack for the resolved mode.

    SOLO mode (default) returns no middleware — the api is a single-user
    local service and bearer tokens would just be noise.

    HOSTED / SHARED modes attach :class:`JWTAuthMiddleware`. In those
    modes ``ATTESTOR_JWT_PUBLIC_KEY`` (PEM or JWKS JSON) MUST be set;
    if it isn't, every request returns 500 ``"auth not configured"`` so
    we fail visibly rather than serving anonymous traffic in a
    multi-tenant deployment.

    Configurable via env (only used when mode != SOLO):

      ATTESTOR_JWT_PUBLIC_KEY     PEM string or path to .pem; required.
      ATTESTOR_JWT_AUDIENCE       Optional ``aud`` claim to enforce.
      ATTESTOR_JWT_ISSUER         Optional ``iss`` claim to enforce.
      ATTESTOR_JWT_ALGORITHMS     Comma-separated; default ``RS256``.
      ATTESTOR_JWT_LEEWAY_SECONDS Default ``0``.
    """
    from starlette.middleware import Middleware

    from attestor.auth import JWTAuthMiddleware
    from attestor.mode import AttestorMode, detect_mode

    mode = detect_mode()
    if mode is AttestorMode.SOLO:
        return []

    public_key = os.environ.get("ATTESTOR_JWT_PUBLIC_KEY")
    if public_key and os.path.isfile(public_key):
        # Allow pointing at a .pem on disk for cloud deploys that
        # mount keys via a volume rather than env-encoding them.
        try:
            with open(public_key, encoding="utf-8") as fh:
                public_key = fh.read()
        except OSError as e:
            logger.error("could not read ATTESTOR_JWT_PUBLIC_KEY=%r: %s",
                         public_key, e)
            public_key = None

    algorithms_raw = os.environ.get("ATTESTOR_JWT_ALGORITHMS", "RS256")
    algorithms = [a.strip() for a in algorithms_raw.split(",") if a.strip()]

    leeway = 0
    try:
        leeway = int(os.environ.get("ATTESTOR_JWT_LEEWAY_SECONDS", "0"))
    except ValueError:
        pass

    return [
        Middleware(
            JWTAuthMiddleware,
            mode=mode,
            public_key=public_key,
            audience=os.environ.get("ATTESTOR_JWT_AUDIENCE"),
            issuer=os.environ.get("ATTESTOR_JWT_ISSUER"),
            algorithms=algorithms,
            leeway_seconds=leeway,
        ),
    ]


app = Starlette(routes=routes, middleware=_build_middleware())
