"""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 os
import logging
import json
import hashlib
from typing import Callable, List, Dict, Any, Optional
from tqdm import tqdm

from pathlib import Path

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

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":
            summary_text = summarizer_instance.summarize_function(path_str, display_name)
        elif symbol_type.upper() == "CLASS":
            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 ThreadPoolExecutor, as_completed, Future

        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)
