"""Tools for indexing and searching LLM-generated code summaries.

This module has two public classes:

* DocstringIndexer – builds a vector index of per-file summaries generated by
  a ``kit.Summarizer``.  The index is stored in a pluggable
  :class:`kit.vector_searcher.VectorDBBackend` (default: Chroma).

* SummarySearcher  – embeds a query and retrieves the most similar summaries
  from a DocstringIndexer backend.
"""

from __future__ import annotations

import hashlib
import json
import logging
import os
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional

from tqdm import tqdm

from .repository import Repository
from .summaries import Summarizer
from .vector_searcher import ChromaDBBackend, VectorDBBackend

EmbedFn = Callable[[str], List[float]]  # str -> embedding vector

__all__ = [
    "DocstringIndexer",
    "SummarySearcher",
]

logger = logging.getLogger(__name__)


def _process_symbol_task(
    path_str: str,
    symbol_info: Dict[str, Any],
    summarizer_instance: Summarizer,
    embed_fn_instance: Callable[[str], List[float]],
    current_cache: Dict[str, Dict[str, str]],  # Pass for reading cache state
) -> Dict[str, Any]:
    """Processes a single symbol: summarize, embed, handle caching and errors."""
    symbol_name = symbol_info.get("name")
    symbol_type = symbol_info.get("type")

    display_name = symbol_info.get("node_path", symbol_name)
    # Ensure doc_id is formed even if symbol_name or display_name is None initially for error reporting
    doc_id_prefix = f"{path_str}::{display_name if display_name else 'unknown_symbol'}"

    if not symbol_name or not symbol_type:
        logger.debug(f"Symbol in {path_str} missing name or type, node_path: {display_name}. Skipping.")
        return {"status": "skipped_no_name_type", "doc_id": doc_id_prefix}

    doc_id = f"{path_str}::{display_name}"  # Re-assign with definite display_name
    symbol_code = symbol_info.get("code", "")

    if not symbol_code:
        logger.warning(f"Could not retrieve code for symbol {display_name} in {path_str}, skipping.")
        return {"status": "skipped_no_code", "doc_id": doc_id}

    try:
        symbol_hash = hashlib.sha1(symbol_code.encode("utf-8", "ignore")).hexdigest()
        if current_cache.get(doc_id, {}).get("hash") == symbol_hash:
            logger.debug(f"Symbol {doc_id} unchanged (hash: {symbol_hash}), skipping.")
            return {"status": "cached", "doc_id": doc_id}

        summary_text = None  # Initialize summary_text
        if symbol_type.upper() == "FUNCTION" or symbol_type.upper() == "METHOD":
            # Ensure display_name is not None before passing
            if display_name is None:
                logger.warning(f"Display name is None for symbol in {path_str}, skipping.")
                return {"status": "skipped_no_display_name", "doc_id": doc_id}
            summary_text = summarizer_instance.summarize_function(path_str, display_name)
        elif symbol_type.upper() == "CLASS":
            # Ensure display_name is not None before passing
            if display_name is None:
                logger.warning(f"Display name is None for symbol in {path_str}, skipping.")
                return {"status": "skipped_no_display_name", "doc_id": doc_id}
            summary_text = summarizer_instance.summarize_class(path_str, display_name)
        else:
            logger.debug(f"Symbol {doc_id} has unsupported type '{symbol_type}'. Skipping summarization.")
            return {"status": "skipped_unsupported_type", "doc_id": doc_id, "symbol_type": symbol_type}

        # Check if summarizer returned a valid string
        if not isinstance(summary_text, str):
            logger.warning(
                f"Summarizer did not return a valid string for symbol {display_name} in {path_str} (type: {symbol_type}). Received: {type(summary_text)}. Skipping."
            )
            return {"status": "skipped_invalid_summary_type", "doc_id": doc_id}

        if not summary_text.strip():
            logger.warning(f"Empty summary for symbol {display_name} in {path_str} (type: {symbol_type}), skipping.")
            return {"status": "skipped_empty_summary", "doc_id": doc_id}

        emb = embed_fn_instance(summary_text)
        meta = {
            "file_path": path_str,
            "symbol_name": display_name,
            "symbol_type": symbol_type,
            "summary": summary_text,
            "level": "symbol",
        }
        return {
            "status": "processed",
            "doc_id": doc_id,
            "embedding": emb,
            "metadata": meta,
            "hash": symbol_hash,
            "summary_for_embedding": summary_text,
        }
    except ValueError as ve:
        logger.warning(f"Skipping symbol {display_name} in {path_str} due to ValueError: {ve}")
        return {"status": "error", "doc_id": doc_id, "error_type": "ValueError", "message": str(ve)}
    except Exception as exc:
        # Log less info from worker to keep main logs cleaner, but still indicate error
        logger.error(f"Failed to process symbol {display_name} in {path_str}: {type(exc).__name__} - {exc}")
        return {"status": "error", "doc_id": doc_id, "error_type": str(type(exc).__name__), "message": str(exc)}


class DocstringIndexer:
    """Builds a vector index of LLM-generated docstrings (summaries).

    A thin wrapper around an existing VectorDB backend.  On ``build()``, it:
    1. walks the repository file tree,
    2. calls :py:meth:`Summarizer.summarize_file` to obtain a concise summary,
    3. embeds that summary via *embed_fn*, and
    4. stores the embedding + metadata in *backend*.

    Parameters
    ----------
    repo
        Active :py:class:`kit.repository.Repository`.
    summarizer
        Configured :py:class:`kit.summaries.Summarizer`.
    embed_fn
        Callable that converts text → embedding vector (list[float]).
    backend
        Optional vector-DB backend, defaults to :class:`ChromaDBBackend`.
    persist_dir
        Where on disk to store backend data (if backend honors persistence).
    """

    def __init__(
        self,
        repo: Repository,
        summarizer: Summarizer,
        embed_fn: Optional[Callable[[str], List[float]]] = None,
        *,
        backend: Optional[VectorDBBackend] = None,
        persist_dir: Optional[str] = None,
    ) -> None:
        self.repo = repo
        self.summarizer = summarizer

        if embed_fn:
            self.embed_fn = embed_fn
        else:
            try:
                from sentence_transformers import SentenceTransformer

                _st_model = SentenceTransformer("all-MiniLM-L6-v2")

                def default_embed_fn(text_to_embed: str) -> List[float]:
                    embedding_vector = _st_model.encode(text_to_embed)
                    return embedding_vector.tolist()

                self.embed_fn = default_embed_fn
                # Example logging:
                # logger.info("No embed_fn provided. Using default SentenceTransformer ('all-MiniLM-L6-v2').")
            except ImportError:
                raise ImportError(
                    "The 'sentence-transformers' library is required to use the default embedding function. "
                    "Please install it (e.g., 'pip install kit[default-embeddings]' or 'pip install sentence-transformers') "
                    "or provide a custom 'embed_fn' to DocstringIndexer."
                )
            except Exception as e:
                raise RuntimeError(
                    f"Failed to initialize the default SentenceTransformer embedding function: {e}. "
                    "Please check your internet connection, ensure the model name is correct, "
                    "or provide a custom 'embed_fn'."
                )

        # Updated persist_dir logic:
        # Default to a .kit_cache directory within the repository path if persist_dir is not given.
        if persist_dir:
            self.persist_dir = persist_dir
        else:
            # Check if repo.repo_path is available and valid, otherwise default to a generic location or raise error
            if hasattr(self.repo, "repo_path") and self.repo.repo_path:
                self.persist_dir = os.path.join(self.repo.repo_path, ".kit_cache", "docstring_db")
            else:
                # Fallback if repo_path is not available (should ideally not happen with a valid Repo object)
                self.persist_dir = os.path.join(os.getcwd(), ".kit_cache", "docstring_db_generic")
                # logger.warning(f"Repository path not found. Defaulting persist_dir to {self.persist_dir}")

        # Ensure the persist_dir exists before ChromaDBBackend tries to use it
        if not os.path.exists(self.persist_dir):
            os.makedirs(self.persist_dir, exist_ok=True)

        # Updated backend instantiation: explicitly pass path and a clearer collection_name
        str_persist_dir = str(self.persist_dir)
        self.backend: VectorDBBackend = backend or ChromaDBBackend(
            persist_dir=str_persist_dir, collection_name="kit_docstring_index"
        )

    def build(self, force: bool = False, level: str = "symbol", file_extensions: Optional[List[str]] = None) -> None:
        """(Re)build the docstring index.

        If *force* is ``False`` and the backend already contains data we do
        nothing.

        Parameters
        ----------
        force
            If ``True``, rebuild even if data exists.
        level
            Granularity of indexing: "file" or "symbol".
        file_extensions
            Optional list of file extensions to include (e.g., [".py", ".js"]).
            If None, all files are considered (respecting .gitignore).
        """
        meta_path = os.path.join(self.persist_dir, "meta.json")
        os.makedirs(self.persist_dir, exist_ok=True)
        if os.path.exists(meta_path):
            with open(meta_path, "r", encoding="utf-8") as fp:
                cache: Dict[str, Dict[str, str]] = json.load(fp)
        else:
            cache = {}

        if force and cache:
            try:
                self.backend.delete(ids=list(cache.keys()))
            except Exception:
                pass
            cache = {}

        all_files = [f["path"] for f in self.repo.get_file_tree() if not f.get("is_dir", False)]

        if file_extensions:
            files_to_process = [fp for fp in all_files if any(fp.endswith(ext) for ext in file_extensions)]
        else:
            files_to_process = all_files

        if not files_to_process:
            logger.warning("No files found to index after extension filtering.")
            return

        embeddings: List[List[float]] = []
        metadatas: List[Dict[str, Any]] = []
        ids: List[str] = []
        seen_ids: set[str] = set()

        from concurrent.futures import Future, ThreadPoolExecutor, as_completed

        try:
            max_workers = int(os.environ.get("KIT_INDEXER_MAX_WORKERS", os.cpu_count() or 4))
        except TypeError:  # os.cpu_count() can return None
            max_workers = 4
        logger.info(f"Using up to {max_workers} workers for symbol processing.")

        for path in tqdm(files_to_process, desc=f"Indexing ({level} level)"):
            if level == "file":
                abs_file_path = Path(self.repo.repo_path) / path  # Create absolute path
                try:
                    summary = self.summarizer.summarize_file(path)
                    if not summary.strip():
                        logger.warning(f"Empty summary for file {path}, skipping.")
                        continue
                    doc_id = path
                    file_hash = hashlib.sha1(abs_file_path.read_bytes()).hexdigest()
                    if cache.get(doc_id, {}).get("hash") == file_hash:
                        seen_ids.add(doc_id)
                        continue
                    meta = {"file_path": path, "summary": summary, "level": "file"}
                    emb = self.embed_fn(summary)
                    embeddings.append(emb)
                    metadatas.append(meta)
                    ids.append(doc_id)
                    cache[doc_id] = {"hash": file_hash}
                    seen_ids.add(doc_id)
                except Exception as exc:
                    logger.error(f"Failed to summarize file {path}: {exc}", exc_info=True)
                    continue
            elif level == "symbol":
                try:
                    symbols_in_file = self.repo.extract_symbols(path)
                except Exception as exc:
                    logger.error(f"Failed to extract symbols from {path}: {exc}", exc_info=True)
                    continue

                if not symbols_in_file:
                    logger.debug(f"No symbols found in {path}, skipping.")
                    continue

                # Lists to store results for the current file
                file_embeddings: List[List[float]] = []
                file_metadatas: List[Dict[str, Any]] = []
                file_ids: List[str] = []
                file_cache_updates: Dict[str, Dict[str, str]] = {}
                file_seen_ids: set[str] = set()

                with ThreadPoolExecutor(max_workers=max_workers) as executor:
                    futures: List[Future] = []
                    for s_info in symbols_in_file:
                        futures.append(
                            executor.submit(
                                _process_symbol_task,
                                path_str=path,
                                symbol_info=s_info,
                                summarizer_instance=self.summarizer,
                                embed_fn_instance=self.embed_fn,
                                current_cache=cache,  # Pass read-only view of cache
                            )
                        )

                    for future in as_completed(futures):
                        try:
                            result = future.result()
                            status = result.get("status")
                            doc_id_from_result = result.get("doc_id", "unknown_doc_id_in_future")

                            if status == "processed":
                                file_embeddings.append(result["embedding"])
                                file_metadatas.append(result["metadata"])
                                file_ids.append(doc_id_from_result)
                                file_cache_updates[doc_id_from_result] = {"hash": result["hash"]}
                                file_seen_ids.add(doc_id_from_result)
                            elif status == "cached":
                                file_seen_ids.add(doc_id_from_result)
                            elif status == "error":
                                logger.debug(
                                    f"Symbol processing for {doc_id_from_result} resulted in status '{status}': {result.get('message')}"
                                )
                            elif status in [
                                "skipped_no_name_type",
                                "skipped_no_code",
                                "skipped_unsupported_type",
                                "skipped_empty_summary",
                                "skipped_invalid_summary_type",
                            ]:
                                logger.debug(f"Symbol {doc_id_from_result} was skipped: {status}")
                            else:
                                logger.warning(
                                    f"Unknown status '{status}' from _process_symbol_task for {doc_id_from_result}"
                                )
                        except Exception as exc_in_future:
                            logger.error(
                                f"Exception retrieving result from future for a symbol in {path}: {exc_in_future}",
                                exc_info=True,
                            )

                # After processing all symbols for the current file, extend the main lists
                if file_embeddings:  # Only extend if there's something to add
                    embeddings.extend(file_embeddings)
                    metadatas.extend(file_metadatas)
                    ids.extend(file_ids)
                    cache.update(file_cache_updates)  # Update the main cache
                    seen_ids.update(file_seen_ids)  # Update main seen_ids
                elif file_seen_ids:  # Still update seen_ids if items were cached
                    seen_ids.update(file_seen_ids)

            else:
                raise ValueError(f"Invalid indexing level: {level}. Choose 'file' or 'symbol'.")

        # remove orphans
        orphan_ids = set(cache.keys()) - seen_ids
        if orphan_ids:
            try:
                self.backend.delete(ids=list(orphan_ids))
            except Exception:
                pass
            for oid in orphan_ids:
                cache.pop(oid, None)

        if embeddings:
            logger.info(f"Adding {len(embeddings)} embeddings to the index.")
            self.backend.add(embeddings=embeddings, metadatas=metadatas, ids=ids)  # Chroma needs 'ids'
            self.backend.persist()
            logger.info("Index build complete and persisted.")
        else:
            logger.warning("No embeddings were generated during indexing.")

        # persist cache
        with open(meta_path, "w", encoding="utf-8") as fp:
            json.dump(cache, fp)

    def get_searcher(self) -> "SummarySearcher":
        """
        Returns a SummarySearcher instance configured with this indexer.

        This provides a convenient way to get a query interface after the
        indexer has been built or loaded.

        Returns:
            SummarySearcher: An instance of SummarySearcher ready to query this indexer.
        """
        return SummarySearcher(indexer=self)


class SummarySearcher:
    """Simple wrapper that queries a :class:`DocstringIndexer` backend."""

    def __init__(self, indexer: DocstringIndexer):
        self.indexer = indexer
        self.embed_fn = indexer.embed_fn

    def search(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
        """Return up to *top_k* hits with their metadata and distance score."""
        if top_k <= 0:
            return []
        emb = self.embed_fn(query)
        return self.indexer.backend.query(emb, top_k)
