"""Documentation indexing module for ChromaDB integration."""

import os
import logging
from pathlib import Path
from typing import Optional, List, Dict
from datetime import datetime, timezone

from .indexing_state import IndexingState
from .git_manager import GitDocumentationManager

logger = logging.getLogger("zk_doc_mcp")

import chromadb


class DocIndexer:
    """Document indexer using ChromaDB for vector search.

    Uses ChromaDB's built-in embeddings (no external model required).
    Requires dependencies: chromadb
    """

    def __init__(
        self,
        doc_path: str,
        collection_name: str = "zk_docs",
        persist_dir: Optional[str] = None,
        force_reindex: bool = False
    ):
        """Initialize the document indexer.

        Args:
            doc_path: Path to documentation directory
            collection_name: Name of the ChromaDB collection
            persist_dir: Directory for persistent storage.
                        Defaults to ~/.zk-doc-mcp/chromadb/
            force_reindex: If True, delete existing index and reindex all docs
        """

        self.doc_path = Path(doc_path)
        self.collection_name = collection_name
        self.force_reindex = force_reindex
        self.state_manager = IndexingState(doc_path)

        if persist_dir is None:
            persist_dir = str(Path.home() / ".zk-doc-mcp" / "chromadb")
        persist_dir = Path(persist_dir)

        # Ensure directory exists
        persist_dir.mkdir(parents=True, exist_ok=True)

        # Initialize PersistentClient
        self.client = chromadb.PersistentClient(path=str(persist_dir))

        # Get or create collection with metadata
        self.collection = self._get_or_create_collection()

    @classmethod
    def from_env(cls):
        """Create DocIndexer instance from environment variables.

        This factory method reads all DocIndexer-related configuration from
        environment variables, providing a single source of truth for default values.

        Environment Variables:
            ZK_DOC_SRC_DIR: Documentation source directory. Default: "~/.zk-doc-mcp/repo"
            ZK_DOC_VECTOR_DB_DIR: Vector database directory. Default: "~/.zk-doc-mcp/chroma_db"
            ZK_DOC_FORCE_REINDEX: Force re-indexing flag. Default: "false"

        Returns:
            DocIndexer instance configured from environment variables

        Example:
            indexer = DocIndexer.from_env()
        """
        return cls(
            doc_path=os.getenv(
                "ZK_DOC_SRC_DIR",
                str(Path.home() / ".zk-doc-mcp" / "repo")
            ),
            persist_dir=os.getenv(
                "ZK_DOC_VECTOR_DB_DIR",
                str(Path.home() / ".zk-doc-mcp" / "chromadb")
            ),
            force_reindex=os.getenv("ZK_DOC_FORCE_REINDEX", "false").lower() == "true"
        )

    def _get_or_create_collection(self):
        """Get existing collection or create new one with metadata."""
        if self.force_reindex:
            # Delete existing collection if exists
            try:
                self.client.delete_collection(name=self.collection_name)
                logger.info(f"Deleted existing collection: {self.collection_name}")
            except ValueError:
                pass  # Collection doesn't exist

        # Create collection with metadata tracking indexing info
        metadata = {
            "schema_version": "1.0",
            "created_at": datetime.now(timezone.utc).isoformat(),
            "last_updated": datetime.now(timezone.utc).isoformat()
        }

        collection = self.client.get_or_create_collection(
            name=self.collection_name,
            metadata=metadata
        )

        return collection

    def _determine_indexing_mode(self) -> tuple[str, str]:
        """Determine which indexing mode should be used based on current state.

        This is the single source of truth for indexing mode selection.
        Returns a tuple of (mode, reason) where mode is one of:
        - "force": Force reindex (delete and rebuild from scratch)
        - "full": Full indexing (no existing index or no indexed commit)
        - "incremental": Incremental update (files have changed since last index)
        - "skip": Skip indexing (already current)

        Returns:
            Tuple of (mode, reason) where mode describes the indexing action
            and reason explains why that mode was selected.
        """
        if self.force_reindex:
            return "force", "Force reindex flag is set"

        if self.is_indexing_in_progress():
            return "force", "Previous indexing was interrupted"

        # Check if collection exists
        try:
            collection_count = self.collection.count()
        except Exception:
            collection_count = 0

        if collection_count == 0:
            return "full", "No existing index found"

        # Collection exists, check if we have indexed commit state
        indexed_commit = self.get_indexed_commit_hash()
        if not indexed_commit:
            return "full", "No indexed commit state found"

        # Check for file changes using git
        try:
            git_manager = GitDocumentationManager(repo_path=str(self.doc_path))
            changed_files = git_manager.get_changed_files(indexed_commit)

            total_changes = len(changed_files.get("added", [])) + \
                          len(changed_files.get("modified", [])) + \
                          len(changed_files.get("deleted", []))

            if total_changes > 0:
                return "incremental", f"Detected {total_changes} changed files"
            else:
                return "skip", "No file changes detected"

        except Exception as e:
            logger.warning(f"Error detecting changes: {e}")
            # Fall back to full reindex on git error
            return "full", "Error detecting changes, falling back to full reindex"

    def index_docs(self, commit_hash: str):
        """Index documentation, choosing between full or incremental based on state.

        Automatically decides indexing mode by analyzing current state:
        - Force mode: If force_reindex flag is set (delete and rebuild)
        - Full mode: If collection doesn't exist or no indexed commit tracked
        - Incremental mode: If collection exists and files have changed
        - Skip: If index is already current

        Args:
            commit_hash: Current git commit hash to index
        """
        try:
            indexing_mode, reason = self._determine_indexing_mode()

            if indexing_mode == "force":
                logger.info(f"Force reindex requested: {reason}")
                self.reset_index()
                self._full_index_docs(commit_hash)
            elif indexing_mode == "full":
                logger.info(f"Full indexing: {reason}")
                self._full_index_docs(commit_hash)
            elif indexing_mode == "incremental":
                logger.info(f"Incremental update: {reason}")
                git_manager = GitDocumentationManager(repo_path=str(self.doc_path))
                indexed_commit = self.get_indexed_commit_hash()
                changed_files = git_manager.get_changed_files(indexed_commit)
                self._incremental_index_docs(changed_files, commit_hash)
            else:  # skip
                logger.info(f"Skipping indexing: {reason}")

        except Exception as e:
            logger.error(f"Error in index_docs: {e}")
            # Fall back to full reindex on error
            logger.warning("Falling back to full reindex due to error")
            self._full_index_docs(commit_hash)

    def _full_index_docs(self, commit_hash: str):
        """Index all markdown files from doc_path (full reindex).

        Args:
            commit_hash: Git commit hash to track this index with
        """
        logger.info(f"Starting full indexing of files from: {self.doc_path}")
        self.state_manager.set_indexing_in_progress()
        indexed_files = 0
        indexed_chunks = 0

        # Batch accumulation for improved performance
        batch_documents = []
        batch_metadatas = []
        batch_ids = []
        batch_size = 500

        for root, _, files in os.walk(self.doc_path):
            for file in files:
                if file.endswith(".md"):
                    file_path = os.path.join(root, file)
                    try:
                        with open(file_path, 'r', encoding='utf-8') as f:
                            content = f.read()
                            chunks = self._chunk_content(content)

                            # Add chunks to batch
                            for i, chunk in enumerate(chunks):
                                batch_documents.append(chunk)
                                batch_metadatas.append({
                                    "source": file_path,
                                    "indexed_at": datetime.now(timezone.utc).isoformat()
                                })
                                batch_ids.append(f"{file_path}-{i}")

                            indexed_files += 1
                            indexed_chunks += len(chunks)

                            # When batch reaches target size, embed and store
                            if len(batch_documents) >= batch_size:
                                self._embed_and_store_batch(
                                    batch_documents,
                                    batch_metadatas,
                                    batch_ids
                                )
                                batch_documents = []
                                batch_metadatas = []
                                batch_ids = []

                    except Exception as e:
                        logger.error(f"Error indexing {file_path}: {e}")

        # Don't forget remaining documents in the batch
        if batch_documents:
            self._embed_and_store_batch(
                batch_documents,
                batch_metadatas,
                batch_ids
            )

        # Update collection metadata
        self.collection.modify(
            metadata={
                **self.collection.metadata,
                "last_updated": datetime.now(timezone.utc).isoformat(),
                "indexed_files": indexed_files,
                "indexed_chunks": indexed_chunks
            }
        )
        self.state_manager.set_indexing_complete(commit_hash)
        logger.info(f"Full indexing complete: {indexed_files} files, {indexed_chunks} chunks")

    def _incremental_index_docs(self, changed_files: dict, commit_hash: str):
        """Incrementally update the index for changed files.

        Only processes files that have been added or modified, and removes
        chunks for deleted files.

        Args:
            changed_files: Dictionary with keys 'added', 'modified', 'deleted'
                          containing lists of file paths
            commit_hash: Git commit hash to update state with
        """
        logger.info("Starting incremental indexing of changed files")
        self.state_manager.set_indexing_in_progress()

        indexed_files = 0
        indexed_chunks = 0
        batch_documents = []
        batch_metadatas = []
        batch_ids = []
        batch_size = 500

        # Process added and modified files
        files_to_process = changed_files.get("added", []) + changed_files.get("modified", [])

        for file_path_rel in files_to_process:
            # Convert relative path to absolute path
            file_path = os.path.join(str(self.doc_path), file_path_rel)

            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    content = f.read()
                    chunks = self._chunk_content(content)

                    # Add chunks to batch
                    for i, chunk in enumerate(chunks):
                        batch_documents.append(chunk)
                        batch_metadatas.append({
                            "source": file_path,
                            "indexed_at": datetime.now(timezone.utc).isoformat()
                        })
                        batch_ids.append(f"{file_path}-{i}")

                    indexed_files += 1
                    indexed_chunks += len(chunks)

                    # When batch reaches target size, upsert (add or update)
                    if len(batch_documents) >= batch_size:
                        self._upsert_batch(
                            batch_documents,
                            batch_metadatas,
                            batch_ids
                        )
                        batch_documents = []
                        batch_metadatas = []
                        batch_ids = []

            except Exception as e:
                logger.error(f"Error processing {file_path}: {e}")

        # Don't forget remaining documents in the batch
        if batch_documents:
            self._upsert_batch(
                batch_documents,
                batch_metadatas,
                batch_ids
            )

        # Delete chunks for removed files using metadata filtering
        for file_path_rel in changed_files.get("deleted", []):
            # Convert relative path to absolute path (same as we did above)
            file_path = os.path.join(str(self.doc_path), file_path_rel)

            try:
                self.collection.delete(
                    where={"source": file_path}
                )
                logger.info(f"Deleted chunks for removed file: {file_path}")
            except Exception as e:
                logger.error(f"Error deleting chunks for {file_path}: {e}")

        # Update collection metadata
        self.collection.modify(
            metadata={
                **self.collection.metadata,
                "last_updated": datetime.now(timezone.utc).isoformat(),
            }
        )

        self.state_manager.set_indexing_complete(commit_hash)
        logger.info(f"Incremental indexing complete: {indexed_files} files, {indexed_chunks} chunks")

    def load_index(self):
        """Load the existing index into memory."""
        # This is implicitly handled by the chromadb client, so we just log it.
        logger.info("Loading existing index.")
        self.get_index_info()

    def get_indexed_commit_hash(self) -> Optional[str]:
        """Get the commit hash of the last successful index."""
        return self.state_manager.get_indexed_commit_hash()

    def is_indexing_in_progress(self) -> bool:
        """Check if an indexing process is currently in progress."""
        return self.state_manager.is_indexing_in_progress()

    def _chunk_content(self, content, chunk_size=1000, chunk_overlap=200):
        chunks = []
        start = 0
        while start < len(content):
            end = start + chunk_size
            chunks.append(content[start:end])
            start += chunk_size - chunk_overlap
        return chunks

    def _embed_and_store_batch(self, documents, metadatas, ids):
        """Store a batch of document chunks with metadata (ChromaDB handles embeddings)."""
        # ChromaDB automatically generates embeddings for documents
        self.collection.add(
            documents=documents,
            metadatas=metadatas,
            ids=ids
        )

    def _upsert_batch(self, documents, metadatas, ids):
        """Upsert a batch of document chunks (add or update if exists).

        Upsert combines add and update operations, allowing us to handle both
        new documents and updated versions of existing documents in one call.
        """
        # ChromaDB automatically generates embeddings for documents
        self.collection.upsert(
            documents=documents,
            metadatas=metadatas,
            ids=ids
        )

    def get_index_info(self) -> dict:
        """Get information about the current index.

        Returns:
            Dictionary with index statistics and metadata
        """
        metadata = self.collection.metadata or {}

        return {
            "collection_name": self.collection_name,
            "schema_version": metadata.get("schema_version"),
            "created_at": metadata.get("created_at"),
            "last_updated": metadata.get("last_updated"),
            "indexed_files": metadata.get("indexed_files"),
            "indexed_chunks": metadata.get("indexed_chunks")
        }

    def reset_index(self):
        """Delete all documents from the collection."""
        self.client.delete_collection(name=self.collection_name)
        logger.info(f"Collection '{self.collection_name}' deleted.")
        self.collection = self._get_or_create_collection()
        logger.info(f"Collection '{self.collection_name}' recreated.")

    def search(self, query: str, n_results: int = 5) -> dict:
        """Search the indexed documents using semantic similarity.

        Args:
            query: Search query string
            n_results: Maximum number of results to return

        Returns:
            Dictionary with search results including documents, distances, and metadata
        """
        # Query the collection (ChromaDB automatically embeds the query)
        results = self.collection.query(
            query_texts=[query],
            n_results=n_results
        )

        # Format results
        formatted_results = []
        if results and results.get('documents'):
            for i, (doc, distance, metadata) in enumerate(
                zip(
                    results['documents'][0],
                    results['distances'][0],
                    results['metadatas'][0]
                )
            ):
                formatted_results.append({
                    "index": i,
                    "document": doc,
                    "distance": float(distance),
                    "metadata": metadata
                })

        return {
            "query": query,
            "results": formatted_results,
            "count": len(formatted_results)
        }

if __name__ == '__main__':
    from .logger import setup_logging
    setup_logging()
    indexer = DocIndexer(doc_path="/Users/hawk/Documents/workspace/DOC/zkdoc/get_started")
    indexer.index_docs("dummy_commit_hash")
    logger.info("Documentation indexed successfully!")
