"""Embed second-brain extractions into the shared public.embeddings table.

This is a near-verbatim copy of the vault-mcp recipe (vault_mcp.core.embeddings),
adapted for our `source_schema='second_brain' / source_table='extractions'`
key tuple and our embed payload (the LLM-generated summary).

Behaviour:
- Reuses the shared `embedding_chunking` library for 512-token chunking with
  64-token overlap.
- Calls Ollama's /api/embeddings with `num_ctx=8192` to embed nomic-embed-text
  vectors (768-d).
- Writes vector literals as `%s::vector` to avoid taking a hard dependency on
  pgvector's Python adapter.
- Delete-before-insert on the (schema, table, source_id, embedding_model)
  tuple so a re-embed of an extraction that previously had N chunks but now
  has M<N doesn't leave orphans.

Graceful degradation: any failure (no DB URL, no Ollama, DB error, empty
summary) logs at WARNING and returns None. The relational `extractions` row
and the file-based vault remain the source of truth — embeddings are an
index, not authoritative state.
"""

from __future__ import annotations

import atexit
import logging
import os
from typing import Any

import httpx
from embedding_chunking import Chunk, chunk_text
from psycopg_pool import ConnectionPool

logger = logging.getLogger(__name__)

SOURCE_SCHEMA = "second_brain"
SOURCE_TABLE = "extractions"

_pool: ConnectionPool | None = None


# ---------------------------------------------------------------------------
# Pool / URL plumbing
# ---------------------------------------------------------------------------


def _database_url() -> str | None:
    """Resolve a raw psycopg DSN (no SQLAlchemy `+psycopg` prefix)."""
    for var in ("SECOND_BRAIN_DATABASE_URL", "HERBYLAB_DATABASE_URL"):
        v = os.environ.get(var)
        if v:
            return v.replace("postgresql+psycopg://", "postgresql://", 1)

    # Fallback: pull from settings.toml so a host-side `process` works even
    # when the operator forgot to export the env var.
    try:
        from second_brain.config import load_config

        url = load_config().database.get("url")
        if url:
            return url.replace("postgresql+psycopg://", "postgresql://", 1)
    except Exception:
        pass
    return None


def _get_pool() -> ConnectionPool | None:
    """Lazy-init a small connection pool. Returns None if no DB URL is configured."""
    global _pool
    if _pool is not None:
        return _pool
    dsn = _database_url()
    if not dsn:
        return None
    _pool = ConnectionPool(dsn, min_size=1, max_size=4, open=True, timeout=10)
    # Make sure the pool's worker threads get a clean shutdown when the
    # interpreter exits; without this psycopg_pool prints
    # "couldn't stop thread pool-1-worker-N within 5.0 seconds" on every
    # `second-brain process` run.
    atexit.register(close_pool)
    return _pool


def close_pool() -> None:
    """Close the embedding-side pool. Safe to call multiple times."""
    global _pool
    if _pool is not None:
        _pool.close()
        _pool = None


# ---------------------------------------------------------------------------
# Embedding model client
# ---------------------------------------------------------------------------


def _ollama_url() -> str:
    """Env var takes precedence over settings.toml so containerized runs can
    point at a different Ollama instance without rebuilding the image."""
    env = os.environ.get("OLLAMA_URL")
    if env:
        return env
    try:
        from second_brain.config import load_config

        return load_config().embeddings.get("ollama_url", "http://ollama:11434")
    except Exception:
        return "http://ollama:11434"


def _embedding_model() -> str:
    env = os.environ.get("EMBEDDING_MODEL")
    if env:
        return env
    try:
        from second_brain.config import load_config

        return load_config().embeddings.get("model", "nomic-embed-text")
    except Exception:
        return "nomic-embed-text"


def embed_text(text: str) -> list[float]:
    """Call Ollama's embeddings endpoint. Returns the model-native vector."""
    resp = httpx.post(
        f"{_ollama_url()}/api/embeddings",
        json={
            "model": _embedding_model(),
            "prompt": text,
            # nomic-embed-text supports 8192-token context. Ollama defaults
            # the runner to 2048, which 500s on long summaries.
            "options": {"num_ctx": 8192},
        },
        timeout=120.0,
    )
    resp.raise_for_status()
    data = resp.json()
    vec = data.get("embedding")
    if not vec or not isinstance(vec, list):
        raise RuntimeError(
            f"Ollama returned no embedding for model {_embedding_model()}"
        )
    return vec


def _vector_literal(vec: list[float]) -> str:
    """pgvector accepts text-literal vectors of form '[v1,v2,...]'."""
    return "[" + ",".join(repr(float(v)) for v in vec) + "]"


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


def embed_extraction(extraction_id: int, summary: str | None) -> int | None:
    """Embed an extraction summary into public.embeddings.

    Args:
        extraction_id: PK of the row in second_brain.extractions.
        summary:       the LLM-generated summary string; we embed this and
                       nothing else for the current round (key_points / claims
                       are not embedded yet).

    Returns the number of chunk rows written, or None on skip/failure.
    Best-effort: failures are logged, not raised.
    """
    if not summary or not summary.strip():
        logger.debug("embed_extraction: skip — extraction %s has no summary", extraction_id)
        return None

    pool = _get_pool()
    if pool is None:
        logger.info(
            "embed_extraction: skip — no SECOND_BRAIN_DATABASE_URL configured"
        )
        return None

    model = _embedding_model()

    try:
        chunks: list[Chunk] = chunk_text(summary)
        chunk_embeddings: list[tuple[Chunk, list[float]]] = [
            (c, embed_text(c.text)) for c in chunks
        ]
    except Exception as exc:
        logger.warning(
            "embed_extraction: chunk/embed failed for extraction %s: %s",
            extraction_id,
            exc,
        )
        return None

    try:
        with pool.connection() as conn:
            n = _upsert_embeddings(
                conn, extraction_id, chunk_embeddings, model
            )
            conn.commit()
            logger.info(
                "embed_extraction: extraction=%s chunks=%d model=%s",
                extraction_id,
                n,
                model,
            )
            return n
    except Exception as exc:
        logger.warning(
            "embed_extraction: DB write failed for extraction %s: %s",
            extraction_id,
            exc,
        )
        return None


def _upsert_embeddings(
    conn: Any,
    extraction_id: int,
    chunk_embeddings: list[tuple[Chunk, list[float]]],
    model: str,
) -> int:
    """Replace all public.embeddings rows for this (extraction, model) pair."""
    with conn.cursor() as cur:
        cur.execute(
            """
            DELETE FROM public.embeddings
             WHERE source_schema = %s
               AND source_table = %s
               AND source_id = %s
               AND embedding_model = %s
            """,
            (SOURCE_SCHEMA, SOURCE_TABLE, extraction_id, model),
        )
        for chunk, vec in chunk_embeddings:
            cur.execute(
                """
                INSERT INTO public.embeddings
                    (source_schema, source_table, source_id, chunk_index,
                     chunk_text, chunk_token_count, embedding, embedding_model)
                VALUES (%s, %s, %s, %s, %s, %s, %s::vector, %s)
                """,
                (
                    SOURCE_SCHEMA,
                    SOURCE_TABLE,
                    extraction_id,
                    chunk.index,
                    chunk.text,
                    chunk.token_count,
                    _vector_literal(vec),
                    model,
                ),
            )
    return len(chunk_embeddings)


__all__ = ["embed_extraction", "embed_text", "close_pool"]