"""Embedding storage — content-SHA-keyed with model tracking, pgvector-ready."""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import uuid

    from sqlalchemy.ext.asyncio import AsyncSession

logger = logging.getLogger(__name__)


class EmbeddingDimensionMismatchError(RuntimeError):
    """Raised when the running embedding provider's dimension does not match
    the vectors already persisted in the ``embeddings`` table.

    Migration 016 replaced the old behavior — which silently ran
    ``UPDATE embeddings SET vector = NULL`` and dropped every stored
    embedding — with a loud refusal. Callers must now explicitly rotate
    to the new dimension via the (phase-4) embeddings rotation endpoint,
    or clear the table manually. Silent data loss is never OK.
    """

    def __init__(self, *, stored: int, expected: int) -> None:
        self.stored = stored
        self.expected = expected
        super().__init__(
            f"embedding dimension mismatch: stored={stored} expected={expected}. "
            f"Refusing to wipe existing vectors. "
            f"Resolve by starting an embedding rotation (POST /embeddings/rotate), "
            f"or explicitly re-dimensioning the column after backing up."
        )


async def ensure_vector_columns(
    session: AsyncSession,
    primary_dim: int,
    shadow_dim: int | None = None,
) -> None:
    """Ensure ``embeddings.vector`` exists at ``primary_dim`` without destroying
    data, and optionally ensure a ``vector_b`` shadow column at ``shadow_dim``.

    Behavior:
      - If ``vector`` does not exist: CREATE at ``primary_dim``, build HNSW index.
      - If ``vector`` exists at ``primary_dim``: no-op.
      - If ``vector`` exists at a DIFFERENT dim: **raise** ``EmbeddingDimensionMismatchError``.
        Old code silently NULLed every row here; that was a footgun.
      - If ``shadow_dim`` is given and differs from ``primary_dim``: ensure a
        ``vector_b`` column at ``shadow_dim`` (used during rotation dual-write).
        Never touches the primary column.

    ``shadow_dim`` is accepted now so the rotation machinery in a later phase
    doesn't have to re-touch this function.
    """
    from sqlalchemy import text

    result = await session.execute(
        text(
            "SELECT atttypmod FROM pg_attribute "
            "WHERE attrelid = 'embeddings'::regclass AND attname = 'vector'"
        )
    )
    row = result.first()

    if row is None:
        # Column doesn't exist — safe to add.
        await session.execute(
            text(f"ALTER TABLE embeddings ADD COLUMN vector vector({primary_dim})")
        )
        await session.execute(
            text(
                "CREATE INDEX IF NOT EXISTS idx_embeddings_vector_hnsw "
                "ON embeddings USING hnsw (vector vector_cosine_ops) "
                "WITH (m = 16, ef_construction = 64)"
            )
        )
        logger.info(
            "Created vector column with dimension %d and HNSW index", primary_dim
        )
    else:
        stored_dim = row[0] - 4  # pgvector stores dim+4 in atttypmod
        if stored_dim != primary_dim:
            logger.error(
                "Embedding dimension mismatch: stored=%d, expected=%d. "
                "Refusing destructive UPDATE. Use embedding rotation flow.",
                stored_dim, primary_dim,
            )
            raise EmbeddingDimensionMismatchError(
                stored=stored_dim, expected=primary_dim,
            )
        logger.debug("Vector column already has correct dimension %d", primary_dim)

    # Shadow column handling — used during a live rotation where new writes
    # target vector_b with the new dimension while reads still hit vector.
    if shadow_dim is not None and shadow_dim != primary_dim:
        result_b = await session.execute(
            text(
                "SELECT atttypmod FROM pg_attribute "
                "WHERE attrelid = 'embeddings'::regclass AND attname = 'vector_b'"
            )
        )
        row_b = result_b.first()
        if row_b is None:
            await session.execute(
                text(f"ALTER TABLE embeddings ADD COLUMN vector_b vector({shadow_dim})")
            )
            await session.execute(
                text(
                    "CREATE INDEX IF NOT EXISTS idx_embeddings_vector_b_hnsw "
                    "ON embeddings USING hnsw (vector_b vector_cosine_ops) "
                    "WITH (m = 16, ef_construction = 64)"
                )
            )
            logger.info("Created shadow vector_b column at dimension %d", shadow_dim)
        else:
            stored_shadow = row_b[0] - 4
            if stored_shadow != shadow_dim:
                raise EmbeddingDimensionMismatchError(
                    stored=stored_shadow, expected=shadow_dim,
                )

    await session.commit()


class EmbeddingStore:
    """Content-SHA-keyed embedding storage for semantic search.

    Embeddings are keyed by (content_sha, embedding_model, embedding_model_version)
    for deduplication. If content_sha+model+version already exists, skip — same
    content always produces the same embeddings for a given model version.

    Branch-aware queries resolve through BranchOverlay.
    """

    def __init__(self, session: AsyncSession) -> None:
        self._session = session

    async def upsert_embeddings(
        self,
        content_sha: str,
        embeddings: list[dict],
    ) -> int:
        """Store embeddings for a content hash.

        Each embedding dict: {chunk_text, chunk_type, embedding_model,
        embedding_model_version?, embedding_dim?, embedding_provenance?, vector?}

        Idempotent within a (content_sha, model, model_version) triple —
        replaces all rows for that triple. Other (model, version) pairs
        coexist peacefully (needed for the dual-write rotation flow).

        Codex fix M7: every insert is mirrored into the ``provenance``
        table so downstream tools can join on ``action_hash`` /
        ``input_blob_oid``. The provenance write is batched alongside
        the embedding insert and does not commit separately — the
        caller's existing session lifecycle controls atomicity.

        Returns count of embeddings stored.
        """
        from sqlalchemy import delete

        from attocode_intel.db.models import Embedding
        from attocode_intel.storage.provenance_store import (
            embedding_provenance_dict,
            write_provenance_rows,
        )

        if not embeddings:
            return 0

        # Determine (model, version) pairs being upserted. Rows at other
        # versions stay put — this is the key invariant enabling rotation.
        pairs_in_batch: set[tuple[str, str]] = {
            (
                e.get("embedding_model", "default"),
                e.get("embedding_model_version", ""),
            )
            for e in embeddings
        }
        for model, version in pairs_in_batch:
            await self._session.execute(
                delete(Embedding).where(
                    Embedding.content_sha == content_sha,
                    Embedding.embedding_model == model,
                    Embedding.embedding_model_version == version,
                )
            )

        # Insert new + collect provenance rows in parallel so they land
        # in the same unit of work.
        provenance_batch: list[dict] = []
        for emb_data in embeddings:
            model = emb_data.get("embedding_model", "default")
            version = emb_data.get("embedding_model_version", "")
            dim = emb_data.get("embedding_dim")
            provenance = emb_data.get("embedding_provenance", {}) or {}
            action_hash = provenance.get("action_hash", "") if isinstance(provenance, dict) else ""
            kwargs = {
                "content_sha": content_sha,
                "embedding_model": model,
                "embedding_model_version": version,
                "embedding_dim": dim,
                "embedding_provenance": provenance,
                "chunk_text": emb_data.get("chunk_text", ""),
                "chunk_type": emb_data.get("chunk_type", "file"),
            }
            # Include vector if provided and pgvector is available
            if "vector" in emb_data and emb_data["vector"] is not None:
                kwargs["vector"] = emb_data["vector"]
            emb = Embedding(**kwargs)
            self._session.add(emb)

            # One provenance row per chunk. Empty action_hash is fine —
            # it just means no dedup key was computed upstream; the row
            # still provides the indexer + blob lineage needed for
            # later audits.
            provenance_batch.append(embedding_provenance_dict(
                action_hash=action_hash,
                content_sha=content_sha,
                model_name=model,
                model_version=version,
                dimension=int(dim) if dim else None,
                config_digest=(
                    provenance.get("config_digest", "")
                    if isinstance(provenance, dict)
                    else ""
                ),
                chunk_type=emb_data.get("chunk_type", "file"),
            ))

        if provenance_batch:
            await write_provenance_rows(self._session, provenance_batch)

        await self._session.flush()
        return len(embeddings)

    async def batch_has_embeddings(self, content_shas: set[str], model: str = "default") -> set[str]:
        """Check which content_shas have embeddings. Returns the subset that do."""
        if not content_shas:
            return set()

        from sqlalchemy import select

        from attocode_intel.db.models import Embedding

        stmt = select(Embedding.content_sha).where(
            Embedding.content_sha.in_(content_shas),
            Embedding.embedding_model == model,
        )
        if hasattr(Embedding, 'vector'):
            stmt = stmt.where(Embedding.vector.isnot(None))
        result = await self._session.execute(stmt.distinct())
        return {row[0] for row in result}

    async def batch_embedding_stats(
        self, content_shas: set[str], model: str = "default",
    ) -> dict[str, dict]:
        """Get chunk count + last embedded time for content SHAs with embeddings."""
        if not content_shas:
            return {}

        from sqlalchemy import func, select

        from attocode_intel.db.models import Embedding

        stmt = select(
            Embedding.content_sha,
            func.count().label("chunk_count"),
            func.max(Embedding.created_at).label("last_embedded"),
        ).where(
            Embedding.content_sha.in_(content_shas),
            Embedding.embedding_model == model,
        )
        if hasattr(Embedding, 'vector'):
            stmt = stmt.where(Embedding.vector.isnot(None))
        stmt = stmt.group_by(Embedding.content_sha)
        result = await self._session.execute(stmt)
        return {
            row.content_sha: {
                "chunk_count": row.chunk_count,
                "last_embedded": row.last_embedded,
            }
            for row in result
        }

    async def has_embeddings(self, content_sha: str, model: str = "default") -> bool:
        """Check if embeddings exist for a content_sha+model. Used for dedup gating."""
        from sqlalchemy import select

        from attocode_intel.db.models import Embedding

        stmt = select(Embedding.id).where(
            Embedding.content_sha == content_sha,
            Embedding.embedding_model == model,
        )
        if hasattr(Embedding, 'vector'):
            stmt = stmt.where(Embedding.vector.isnot(None))
        result = await self._session.execute(stmt.limit(1))
        return result.scalar_one_or_none() is not None

    async def find_similar_by_sha(
        self,
        branch_id: uuid.UUID,
        content_sha: str,
        top_k: int = 10,
        model: str = "default",
    ) -> list[dict]:
        """Find files similar to the given content_sha within a branch.

        Retrieves the embedding vector(s) for the source content_sha, averages
        them if multiple chunks exist, then runs a cosine similarity search
        excluding the source itself.
        """
        from sqlalchemy import text

        from attocode_intel.storage.branch_overlay import BranchOverlay

        # 1. Get vectors for the source content_sha
        result = await self._session.execute(
            text("""
                SELECT vector::text FROM embeddings
                WHERE content_sha = :sha AND embedding_model = :model AND vector IS NOT NULL
            """),
            {"sha": content_sha, "model": model},
        )
        rows = result.fetchall()
        if not rows:
            return []

        # 2. Parse and average vectors
        def parse_vector(vec_str: str) -> list[float]:
            return [float(x) for x in vec_str.strip("[]").split(",")]

        vectors = [parse_vector(row[0]) for row in rows]
        dim = len(vectors[0])
        if len(vectors) == 1:
            avg_vector = vectors[0]
        else:
            avg_vector = [sum(v[i] for v in vectors) / len(vectors) for i in range(dim)]

        # 3. Resolve branch manifest
        overlay = BranchOverlay(self._session)
        manifest = await overlay.resolve_manifest(branch_id)
        content_shas = set(manifest.values())
        content_shas.discard(content_sha)  # exclude source

        if not content_shas:
            return []

        sha_to_path = {sha: path for path, sha in manifest.items()}

        # 4. Run cosine similarity search
        result = await self._session.execute(
            text("""
                SELECT content_sha, chunk_text, chunk_type, embedding_model,
                       1 - (vector <=> CAST(:qv AS vector)) AS score
                FROM embeddings
                WHERE content_sha = ANY(:shas)
                  AND embedding_model = :model
                  AND vector IS NOT NULL
                ORDER BY vector <=> CAST(:qv AS vector)
                LIMIT :top_k
            """),
            {
                "qv": "[" + ",".join(str(v) for v in avg_vector) + "]",
                "shas": list(content_shas),
                "model": model,
                "top_k": top_k,
            },
        )

        results = []
        for row in result:
            results.append({
                "file": sha_to_path.get(row.content_sha, "unknown"),
                "content_sha": row.content_sha,
                "chunk_text": row.chunk_text,
                "score": float(row.score),
            })
        return results

    async def similarity_search(
        self,
        branch_id: uuid.UUID,
        query_vector: list[float],
        top_k: int = 10,
        model: str = "default",
        file_filter: str = "",
    ) -> list[dict]:
        """Find most similar content within a branch context using pgvector cosine distance.

        Args:
            branch_id: Branch to scope results to.
            query_vector: Embedded query vector.
            top_k: Number of results to return.
            model: Embedding model name to filter by.

        Returns:
            List of dicts with file, content_sha, chunk_text, chunk_type, model, score.
        """
        from sqlalchemy import text

        from attocode_intel.storage.branch_overlay import BranchOverlay

        overlay = BranchOverlay(self._session)
        manifest = await overlay.resolve_manifest(branch_id)

        if file_filter:
            import fnmatch
            manifest = {p: s for p, s in manifest.items() if fnmatch.fnmatch(p, file_filter)}

        content_shas = set(manifest.values())

        if not content_shas:
            return []

        sha_to_path = {sha: path for path, sha in manifest.items()}

        # pgvector cosine distance query — score = 1 - distance (higher = more similar)
        result = await self._session.execute(
            text("""
                SELECT content_sha, chunk_text, chunk_type, embedding_model,
                       1 - (vector <=> CAST(:qv AS vector)) AS score
                FROM embeddings
                WHERE content_sha = ANY(:shas)
                  AND embedding_model = :model
                  AND vector IS NOT NULL
                ORDER BY vector <=> CAST(:qv AS vector)
                LIMIT :top_k
            """),
            {
                "qv": "[" + ",".join(str(v) for v in query_vector) + "]",
                "shas": list(content_shas),
                "model": model,
                "top_k": top_k,
            },
        )

        results = []
        for row in result:
            results.append({
                "file": sha_to_path.get(row.content_sha, "unknown"),
                "content_sha": row.content_sha,
                "chunk_text": row.chunk_text,
                "chunk_type": row.chunk_type,
                "model": row.embedding_model,
                "score": float(row.score),
            })
        return results

    async def multi_branch_similarity_search(
        self,
        session: AsyncSession,
        branch_manifests: dict[str, dict[str, str]],
        query_vector: list[float],
        top_k: int = 20,
        model: str = "",
        file_filter: str = "",
    ) -> list[dict]:
        """Find most similar content across multiple repo branch manifests.

        Args:
            session: Async DB session (uses self._session if same, but kept
                     for API symmetry with callers that pass it explicitly).
            branch_manifests: Mapping of repo_id (str) → manifest (path → content_sha).
            query_vector: Embedded query vector.
            top_k: Number of results to return globally.
            model: Embedding model name to filter by.
            file_filter: Optional glob pattern to filter file paths.

        Returns:
            List of dicts with repo_id, file, content_sha, chunk_text,
            chunk_type, model, score — sorted by score descending.
        """
        import fnmatch

        from sqlalchemy import text

        # 1. Union all content_shas across repos, building reverse lookup
        all_content_shas: set[str] = set()
        # sha → (repo_id, path) for mapping results back
        sha_to_context: dict[str, tuple[str, str]] = {}

        for repo_id, manifest in branch_manifests.items():
            for path, sha in manifest.items():
                if file_filter and not fnmatch.fnmatch(path, file_filter):
                    continue
                all_content_shas.add(sha)
                # First repo wins for a given sha (content-addressed dedup)
                if sha not in sha_to_context:
                    sha_to_context[sha] = (repo_id, path)

        if not all_content_shas:
            return []

        # 2. Single pgvector query across all repos
        result = await self._session.execute(
            text("""
                SELECT content_sha, chunk_text, chunk_type, embedding_model,
                       1 - (vector <=> CAST(:qv AS vector)) AS score
                FROM embeddings
                WHERE content_sha = ANY(:shas)
                  AND embedding_model = :model
                  AND vector IS NOT NULL
                ORDER BY vector <=> CAST(:qv AS vector)
                LIMIT :top_k
            """),
            {
                "qv": "[" + ",".join(str(v) for v in query_vector) + "]",
                "shas": list(all_content_shas),
                "model": model,
                "top_k": top_k,
            },
        )

        # 3. Map results back to repo context
        results = []
        for row in result:
            ctx = sha_to_context.get(row.content_sha)
            repo_id = ctx[0] if ctx else "unknown"
            file_path = ctx[1] if ctx else "unknown"
            results.append({
                "repo_id": repo_id,
                "file": file_path,
                "content_sha": row.content_sha,
                "chunk_text": row.chunk_text,
                "chunk_type": row.chunk_type,
                "model": row.embedding_model,
                "score": float(row.score),
            })
        return results

    async def gc_orphaned(
        self,
        min_age_minutes: int = 60,
        *,
        repo_id: uuid.UUID | None = None,
    ) -> int:
        """Delete embeddings whose content_sha is not in any branch manifest.

        Only deletes embeddings older than *min_age_minutes* to avoid racing
        with concurrent indexers that are still writing references.

        When ``repo_id`` is provided, restricts the DELETE to embeddings
        whose content_sha was ever tracked by a branch of the target
        repo AND is no longer tracked anywhere. ``repo_id=None`` runs
        global GC (used by background workers).

        Returns count of deleted rows.
        """
        from sqlalchemy import text

        if repo_id is None:
            result = await self._session.execute(
                text("""
                DELETE FROM embeddings
                WHERE created_at < NOW() - make_interval(mins => :age)
                  AND content_sha NOT IN (
                    SELECT DISTINCT content_sha FROM branch_files WHERE content_sha IS NOT NULL
                )
                """).bindparams(age=min_age_minutes)
            )
        else:
            result = await self._session.execute(
                text("""
                DELETE FROM embeddings
                WHERE created_at < NOW() - make_interval(mins => :age)
                  AND content_sha IN (
                    SELECT DISTINCT bf.content_sha
                    FROM branch_files bf
                    JOIN branches b ON b.id = bf.branch_id
                    WHERE b.repo_id = :repo_id AND bf.content_sha IS NOT NULL
                  )
                  AND content_sha NOT IN (
                    SELECT DISTINCT content_sha FROM branch_files WHERE content_sha IS NOT NULL
                )
                """).bindparams(age=min_age_minutes, repo_id=str(repo_id))
            )
        count = result.rowcount
        if count:
            if repo_id is None:
                logger.info(
                    "GC: removed %d orphaned embeddings (older than %dm, global)",
                    count, min_age_minutes,
                )
            else:
                logger.info(
                    "GC: removed %d orphaned embeddings (older than %dm, repo=%s)",
                    count, min_age_minutes, repo_id,
                )
        return count
