"""ChromaDB vector store — embedded mode, domain-oriented collections."""

from __future__ import annotations

import logging
import re
import sys
from pathlib import Path
from typing import Any, cast

# Workaround: chromadb uses pydantic v1 BaseSettings which fails on Python 3.14+
# due to PEP 649 lazy annotations. Patch pydantic v1 to tolerate unknown field types.
if sys.version_info >= (3, 14):
    import pydantic.v1.fields as _pv1_fields

    _orig_set_default_and_type = _pv1_fields.ModelField._set_default_and_type

    def _patched_set_default_and_type(self: Any) -> None:
        try:
            _orig_set_default_and_type(self)
        except _pv1_fields.errors_.ConfigError:
            self.outer_type_ = self.type_ = Any
            self.required = False

    _pv1_fields.ModelField._set_default_and_type = _patched_set_default_and_type

import chromadb

from redhat_oracle.core.config import OracleConfig, get_data_dir
from redhat_oracle.core.vectordb.chunker import Chunk
from redhat_oracle.core.vectordb.embeddings import EmbeddingProvider

logger = logging.getLogger(__name__)

COLLECTION_NAMES = ("oracle_code", "oracle_jira", "oracle_slack", "oracle_web", "oracle_pages")

SOURCE_TO_COLLECTION = {
    "github": "oracle_code",
    "gitlab": "oracle_code",
    "jira": "oracle_jira",
    "slack": "oracle_slack",
    "web": "oracle_web",
    "pages": "oracle_pages",
}

_OLD_COLLECTION_NAMES = ("oracle_github", "oracle_gitlab")


class VectorStore:
    """ChromaDB-backed vector store with isolated databases per collection.

    Each collection gets its own ChromaDB PersistentClient and sqlite file,
    so corruption in one collection cannot affect the others.
    """

    def __init__(
        self,
        config: OracleConfig,
        embedder: EmbeddingProvider | None = None,
        db_path: Path | None = None,
    ) -> None:
        self._config = config
        self._embedder = embedder or EmbeddingProvider(
            model_name=config.vectordb.embeddings.model,
        )
        self._batch_size = config.vectordb.embeddings.batch_size

        base_dir = db_path or (get_data_dir(config) / "vectordb" / "chroma")
        self._clients: dict[str, Any] = {}
        self._collections: dict[str, chromadb.Collection] = {}

        for name in COLLECTION_NAMES:
            col_dir = base_dir / name
            col_dir.mkdir(parents=True, exist_ok=True)
            client = chromadb.PersistentClient(path=str(col_dir))
            self._clients[name] = client
            self._collections[name] = client.get_or_create_collection(
                name=name,
                metadata={"hnsw:space": "cosine"},
            )

    def _get_collection(self, source: str) -> chromadb.Collection:
        """Get the collection for a source type."""
        collection_name = SOURCE_TO_COLLECTION.get(source)
        if collection_name is None or collection_name not in self._collections:
            raise ValueError(f"Unknown source type: {source}")
        return self._collections[collection_name]

    def upsert_chunks(
        self,
        chunks: list[Chunk],
        progress_callback: Any | None = None,
    ) -> int:
        """Upsert chunks into the appropriate collection, batched.

        Args:
            chunks: List of chunks to upsert.
            progress_callback: Optional callable(batch_num, total_batches) for progress.

        Returns:
            Number of chunks upserted.
        """
        if not chunks:
            return 0

        # Deduplicate chunks by ID (keep last occurrence)
        seen: dict[str, Chunk] = {}
        for chunk in chunks:
            seen[chunk.chunk_id] = chunk
        deduped = list(seen.values())

        if len(deduped) < len(chunks):
            logger.info(
                "Deduplicated %d chunks down to %d",
                len(chunks),
                len(deduped),
            )

        # Group chunks by source
        by_source: dict[str, list[Chunk]] = {}
        for chunk in deduped:
            by_source.setdefault(chunk.source, []).append(chunk)

        total_upserted = 0

        for source, source_chunks in by_source.items():
            collection = self._get_collection(source)
            total_batches = (len(source_chunks) + self._batch_size - 1) // self._batch_size

            for batch_idx in range(total_batches):
                start = batch_idx * self._batch_size
                end = start + self._batch_size
                batch = source_chunks[start:end]

                texts = [c.text for c in batch]
                ids = [c.chunk_id for c in batch]
                metadatas = [
                    {
                        "source": c.source,
                        "org": c.org,
                        "repo": c.repo,
                        "path": c.path,
                        "start_line": c.start_line,
                        "end_line": c.end_line,
                        "node_type": c.node_type or "",
                        "node_name": c.node_name or "",
                        "author": c.author or "",
                        "timestamp": c.timestamp or "",
                        "url": c.url or "",
                        "jira_issue_id": c.jira_issue_id or "",
                        "content_type": c.content_type or "",
                        "status": c.status or "",
                        "priority": c.priority or "",
                        "labels": ",".join(c.labels) if c.labels else "",
                        "pr_number": c.pr_number or "",
                    }
                    for c in batch
                ]

                embeddings = self._embedder.embed(texts)

                collection.upsert(
                    ids=ids,
                    embeddings=cast(Any, embeddings),
                    documents=texts,
                    metadatas=cast(Any, metadatas),
                )

                total_upserted += len(batch)

                if progress_callback:
                    progress_callback(batch_idx + 1, total_batches)

        return total_upserted

    def delete_repo_chunks(self, source: str, org: str, repo: str) -> int:
        """Delete all chunks for a specific repo (for full re-index).

        Returns:
            Number of chunks deleted.
        """
        collection = self._get_collection(source)

        # Query for all chunks belonging to this repo
        results = collection.get(
            where={"$and": [{"org": org}, {"repo": repo}]},
        )

        if not results["ids"]:
            return 0

        count = len(results["ids"])
        collection.delete(ids=results["ids"])
        logger.info("Deleted %d chunks for %s/%s/%s", count, source, org, repo)
        return count

    def reindex_repo(
        self,
        chunks: list[Chunk],
        source: str,
        org: str,
        repo: str,
        progress_callback: Any | None = None,
    ) -> dict[str, int]:
        """Full re-index: delete existing chunks for a repo, then insert new ones.

        Returns:
            Dict with 'deleted' and 'inserted' counts.
        """
        deleted = self.delete_repo_chunks(source, org, repo)
        inserted = self.upsert_chunks(chunks, progress_callback)
        return {"deleted": deleted, "inserted": inserted}

    # Pattern for Jira-style issue keys: PROJECT-123
    _JIRA_KEY_RE = re.compile(r"\b([A-Z][A-Z0-9]+-\d+)\b")

    def _lookup_by_metadata(
        self,
        collection: chromadb.Collection,
        field: str,
        value: str,
    ) -> list[dict[str, Any]]:
        """Fetch chunks by exact metadata match, returned with score 1.0."""
        try:
            results = collection.get(
                where={field: value},
                include=["documents", "metadatas"],
            )
        except Exception:
            return []

        items: list[dict[str, Any]] = []
        for i, doc_id in enumerate(results["ids"]):
            items.append({
                "id": doc_id,
                "text": results["documents"][i] if results["documents"] else "",
                "metadata": results["metadatas"][i] if results["metadatas"] else {},
                "score": 1.0,
            })
        return items

    def search_by_metadata(
        self,
        field: str,
        value: str,
        source: str | None = None,
        top_k: int = 5,
    ) -> list[dict[str, Any]]:
        """Search across collections by exact metadata field match.

        Args:
            field: Metadata field name (e.g. "author", "jira_issue_id").
            value: Exact value to match.
            source: Optional source type filter to search a single collection.
            top_k: Maximum number of results.

        Returns:
            List of result dicts with 'id', 'text', 'metadata', 'score' (always 1.0).
        """
        if not value:
            return []

        collections = [self._get_collection(source)] if source else list(self._collections.values())

        all_results: list[dict[str, Any]] = []
        for col in collections:
            items = self._lookup_by_metadata(col, field, value)
            all_results.extend(items)
            if len(all_results) >= top_k:
                break

        return all_results[:top_k]

    def search(
        self,
        query: str,
        source: str | None = None,
        top_k: int = 10,
        min_score: float = 0.3,
    ) -> list[dict[str, Any]]:
        """Search for similar chunks.

        Combines semantic vector search with exact metadata lookups for
        identifiers like Jira issue keys (e.g., AAP-1234).

        Args:
            query: The search query text.
            source: Optional source type filter ("github", "gitlab", "slack").
            top_k: Maximum number of results.
            min_score: Minimum similarity score (cosine).

        Returns:
            List of result dicts with 'text', 'metadata', 'score'.
        """
        all_results: list[dict[str, Any]] = []
        seen_ids: set[str] = set()

        # --- Exact metadata lookups for known patterns ---
        jira_keys = self._JIRA_KEY_RE.findall(query)
        if jira_keys:
            jira_col = self._collections.get("oracle_jira")
            if jira_col and (source is None or source == "jira"):
                for key in jira_keys:
                    for item in self._lookup_by_metadata(
                        jira_col, "jira_issue_id", key
                    ):
                        if item["id"] not in seen_ids:
                            seen_ids.add(item["id"])
                            all_results.append(item)

        # --- Semantic vector search ---
        query_embedding = self._embedder.embed_single(query)

        collections_to_search: list[chromadb.Collection]
        if source:
            collections_to_search = [self._get_collection(source)]
        else:
            collections_to_search = list(self._collections.values())

        for collection in collections_to_search:
            try:
                count = collection.count()
            except Exception:
                count = 0

            if count == 0:
                continue

            results = collection.query(
                query_embeddings=cast(Any, [query_embedding]),
                n_results=min(top_k, count),
                include=["documents", "metadatas", "distances"],
            )

            if not results["ids"] or not results["ids"][0]:
                continue

            for i, doc_id in enumerate(results["ids"][0]):
                if doc_id in seen_ids:
                    continue

                # ChromaDB cosine distance: 0 = identical, 2 = opposite
                # Convert to similarity: 1 - (distance / 2)
                distance = results["distances"][0][i] if results["distances"] else 0.0
                score = 1.0 - (distance / 2.0)

                if score < min_score:
                    continue

                seen_ids.add(doc_id)
                all_results.append({
                    "id": doc_id,
                    "text": results["documents"][0][i] if results["documents"] else "",
                    "metadata": results["metadatas"][0][i] if results["metadatas"] else {},
                    "score": score,
                })

        # Sort by score descending and return top-k
        all_results.sort(key=lambda r: r["score"], reverse=True)
        return all_results[:top_k]

    def collection_stats(self) -> dict[str, int]:
        """Return chunk counts per collection."""
        stats: dict[str, int] = {}
        for name, col in self._collections.items():
            try:
                stats[name] = col.count()
            except Exception:
                logger.warning("Failed to read collection %s (possibly corrupted)", name)
                stats[name] = -1
        return stats

    def migrate_collections(self) -> dict[str, int]:
        """Migrate data from old oracle_github/oracle_gitlab collections.

        No longer needed with separate databases — kept for API compatibility.
        """
        return {name: 0 for name in _OLD_COLLECTION_NAMES}
