"""
Vector Store for semantic search in Natural Language Query system.

Uses ChromaDB for storing and retrieving:
- Schema metadata embeddings
- Example query embeddings
- Business glossary embeddings
- Cached successful queries
"""

from typing import List, Dict, Any, Optional, Callable
import os
import logging
import time
from pathlib import Path
from functools import wraps
import pandas as pd

# Configure logging
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)


def retry_on_failure(max_retries: int = 3, delay: float = 0.5):
    """
    Decorator to retry ChromaDB operations on failure.

    Args:
        max_retries: Maximum number of retry attempts
        delay: Delay in seconds between retries (exponential backoff)
    """
    def decorator(func: Callable) -> Callable:
        @wraps(func)
        def wrapper(*args, **kwargs):
            last_exception = None
            for attempt in range(max_retries):
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    last_exception = e
                    if attempt < max_retries - 1:
                        wait_time = delay * (2 ** attempt)  # Exponential backoff
                        logger.warning(
                            f"{func.__name__} failed (attempt {attempt + 1}/{max_retries}): {e}. "
                            f"Retrying in {wait_time:.1f}s..."
                        )
                        time.sleep(wait_time)
                    else:
                        logger.error(
                            f"{func.__name__} failed after {max_retries} attempts: {e}"
                        )
            raise last_exception
        return wrapper
    return decorator


class VectorStore:
    """
    Vector database for storing and retrieving relevant code examples.

    Uses ChromaDB with sentence-transformers embeddings for semantic search.
    Stores example queries, code snippets, and documentation.
    """

    def __init__(
        self,
        collection_name: str = "outhad_edge_nlq",
        persist_directory: Optional[str] = None,
        embedding_model: str = "all-MiniLM-L6-v2",
    ):
        """
        Initialize vector store.

        Args:
            collection_name: Name of the ChromaDB collection
            persist_directory: Directory to persist the database
            embedding_model: Sentence transformer model name
        """
        try:
            import chromadb
            from chromadb.config import Settings
        except ImportError:
            raise ImportError(
                "chromadb is required for vector store. "
                "Install with: pip install outhad_edge[ai]"
            )

        self.collection_name = collection_name
        self.embedding_model = embedding_model

        # Setup persist directory
        if persist_directory is None:
            persist_directory = os.getenv(
                "CHROMA_DB_PATH", str(Path.home() / ".outhad_edge" / "chromadb")
            )

        self.persist_directory = persist_directory
        Path(self.persist_directory).mkdir(parents=True, exist_ok=True)

        # Initialize ChromaDB client
        self.client = chromadb.Client(
            Settings(
                persist_directory=self.persist_directory,
                anonymized_telemetry=False,
            )
        )

        # Get or create main collection (for backward compatibility)
        try:
            self.collection = self.client.get_collection(name=collection_name)
        except:
            self.collection = self.client.create_collection(
                name=collection_name,
                metadata={"description": "Outhad_Edge NLQ examples and documentation"},
            )

        # Create additional collections for comprehensive RAG
        try:
            self.schema_collection = self.client.get_or_create_collection(
                name="schema_metadata",
                metadata={"description": "Table and column metadata"}
            )
        except:
            self.schema_collection = None

        try:
            self.examples_collection = self.client.get_or_create_collection(
                name="example_queries",
                metadata={"description": "Example NL query → code pairs"}
            )
        except:
            self.examples_collection = None

        try:
            self.glossary_collection = self.client.get_or_create_collection(
                name="business_glossary",
                metadata={"description": "Business terms and definitions"}
            )
        except:
            self.glossary_collection = None

        try:
            self.cache_collection = self.client.get_or_create_collection(
                name="query_cache",
                metadata={"description": "Cached successful queries"}
            )
        except:
            self.cache_collection = None

    @retry_on_failure(max_retries=3, delay=0.5)
    def add_examples(self, examples: List[Dict[str, Any]]) -> None:
        """
        Add code examples to the vector store (backward compatibility method).

        Supports both 'question' and 'query' keys for backward compatibility.

        Args:
            examples: List of examples with 'question'/'query', 'code', and 'description'

        Raises:
            ValueError: If examples are invalid or missing required fields
            RuntimeError: If ChromaDB operation fails after retries
        """
        if not examples:
            logger.warning("add_examples called with empty examples list")
            return

        if not isinstance(examples, list):
            raise ValueError(f"examples must be a list, got {type(examples)}")

        documents = []
        metadatas = []
        ids = []

        for i, example in enumerate(examples):
            if not isinstance(example, dict):
                logger.error(f"Example {i} is not a dict: {type(example)}")
                raise ValueError(f"Example {i} must be a dict, got {type(example)}")

            # Support both 'query' and 'question' keys for backward compatibility
            question = example.get('question') or example.get('query')
            if not question:
                logger.error(f"Example {i} missing both 'question' and 'query' keys: {example.keys()}")
                raise ValueError(
                    f"Example {i} must have either 'question' or 'query' key. "
                    f"Found keys: {list(example.keys())}"
                )

            code = example.get('code')
            if not code:
                logger.error(f"Example {i} missing 'code' key: {example.keys()}")
                raise ValueError(f"Example {i} must have 'code' key")

            # Combine question and description for better semantic search
            description = example.get('description', '')
            doc_text = f"{question} {description}".strip()
            documents.append(doc_text)

            metadatas.append({
                "question": question,  # Normalize to 'question' in storage
                "code": code,
                "description": description,
                "tags": ",".join(example.get("tags", [])),
            })

            ids.append(f"example_{i}_{abs(hash(question))}")

        try:
            self.collection.add(documents=documents, metadatas=metadatas, ids=ids)
            logger.info(f"Successfully added {len(examples)} examples to main collection")
        except Exception as e:
            # Try to handle duplicate IDs
            logger.warning(f"Failed to add examples, attempting to handle duplicates: {e}")
            try:
                self.collection.delete(ids=ids)
                self.collection.add(documents=documents, metadatas=metadatas, ids=ids)
                logger.info(f"Successfully added {len(examples)} examples after handling duplicates")
            except Exception as e2:
                logger.error(f"Failed to add examples even after duplicate handling: {e2}")
                raise RuntimeError(f"Failed to add examples to vector store: {e2}") from e2

    def search(
        self, query: str, n_results: int = 3, filter_tags: Optional[List[str]] = None
    ) -> List[Dict[str, Any]]:
        """
        Search for relevant examples based on semantic similarity.

        Args:
            query: User's natural language query
            n_results: Number of results to return
            filter_tags: Optional tags to filter by

        Returns:
            list: Relevant examples with code and descriptions
        """
        # Build query filter
        where = None
        if filter_tags:
            where = {"tags": {"$in": filter_tags}}

        # Search
        results = self.collection.query(
            query_texts=[query], n_results=n_results, where=where
        )

        # Format results
        examples = []
        if results["metadatas"] and results["metadatas"][0]:
            for metadata in results["metadatas"][0]:
                examples.append(
                    {
                        "question": metadata["question"],
                        "code": metadata["code"],
                        "description": metadata["description"],
                        "tags": metadata.get("tags", "").split(",")
                        if metadata.get("tags")
                        else [],
                    }
                )

        return examples

    def add_documentation(self, doc_text: str, metadata: Dict[str, Any]) -> None:
        """
        Add documentation snippet to vector store.

        Args:
            doc_text: Documentation text
            metadata: Metadata about the documentation
        """
        doc_id = f"doc_{hash(doc_text)}"

        self.collection.add(
            documents=[doc_text], metadatas=[metadata], ids=[doc_id]
        )

    def clear(self) -> None:
        """Clear all data from the collection."""
        self.client.delete_collection(name=self.collection_name)
        self.collection = self.client.create_collection(
            name=self.collection_name,
            metadata={"description": "Outhad_Edge NLQ examples and documentation"},
        )

    def count(self) -> int:
        """Get number of items in the collection."""
        return self.collection.count()

    def get_all_examples(self) -> List[Dict[str, Any]]:
        """
        Get all examples from the vector store.

        Returns:
            list: All stored examples
        """
        results = self.collection.get()

        examples = []
        if results["metadatas"]:
            for metadata in results["metadatas"]:
                examples.append(
                    {
                        "question": metadata["question"],
                        "code": metadata["code"],
                        "description": metadata["description"],
                        "tags": metadata.get("tags", "").split(",")
                        if metadata.get("tags")
                        else [],
                    }
                )

        return examples

    @retry_on_failure(max_retries=3, delay=0.5)
    def add_schema_metadata(self, semantic_layer) -> None:
        """
        Add schema metadata to vector store.

        Args:
            semantic_layer: SemanticLayer instance with schema information

        Raises:
            ValueError: If semantic_layer is invalid
            RuntimeError: If ChromaDB operation fails after retries
        """
        if not self.schema_collection:
            logger.warning("schema_collection not initialized, skipping add_schema_metadata")
            return

        if semantic_layer is None:
            logger.error("semantic_layer is None")
            raise ValueError("semantic_layer cannot be None")

        try:
            schema_desc = semantic_layer.get_schema_description()
            api_doc = semantic_layer.get_eventstream_api_doc()
        except Exception as e:
            logger.error(f"Failed to extract schema metadata from semantic_layer: {e}")
            raise ValueError(f"Failed to extract schema metadata: {e}") from e

        if not schema_desc or not api_doc:
            logger.error(f"Empty schema metadata: schema_desc={bool(schema_desc)}, api_doc={bool(api_doc)}")
            raise ValueError("Schema description or API doc is empty")

        # Add to collection
        try:
            self.schema_collection.add(
                documents=[schema_desc, api_doc],
                ids=["eventstream_schema", "eventstream_api"],
                metadatas=[
                    {"type": "schema", "table": "eventstream"},
                    {"type": "api_doc", "table": "eventstream"}
                ]
            )
            logger.info("Successfully added schema metadata (2 documents)")
        except Exception as e:
            # Handle duplicate IDs
            logger.warning(f"Failed to add schema metadata, attempting to handle duplicates: {e}")
            try:
                self.schema_collection.delete(ids=["eventstream_schema", "eventstream_api"])
                self.schema_collection.add(
                    documents=[schema_desc, api_doc],
                    ids=["eventstream_schema", "eventstream_api"],
                    metadatas=[
                        {"type": "schema", "table": "eventstream"},
                        {"type": "api_doc", "table": "eventstream"}
                    ]
                )
                logger.info("Successfully added schema metadata after handling duplicates")
            except Exception as e2:
                logger.error(f"Failed to add schema metadata even after duplicate handling: {e2}")
                raise RuntimeError(f"Failed to add schema metadata to vector store: {e2}") from e2

    @retry_on_failure(max_retries=3, delay=0.5)
    def add_example_queries(self, examples: List[Dict[str, str]]) -> None:
        """
        Add example query → code pairs to vector store.

        Supports both 'query' and 'question' keys for backward compatibility.

        Args:
            examples: List of examples with 'query'/'question', 'code', 'answer'

        Raises:
            ValueError: If examples are invalid or missing required fields
            RuntimeError: If ChromaDB operation fails after retries
        """
        if not self.examples_collection:
            logger.warning("examples_collection not initialized, skipping add_example_queries")
            return

        if not examples:
            logger.warning("add_example_queries called with empty examples list")
            return

        if not isinstance(examples, list):
            raise ValueError(f"examples must be a list, got {type(examples)}")

        documents = []
        ids = []
        metadatas = []

        for i, example in enumerate(examples):
            if not isinstance(example, dict):
                logger.error(f"Example {i} is not a dict: {type(example)}")
                raise ValueError(f"Example {i} must be a dict, got {type(example)}")

            # Support both 'query' and 'question' keys for backward compatibility
            query = example.get('query') or example.get('question')
            if not query:
                logger.error(f"Example {i} missing both 'query' and 'question' keys: {example.keys()}")
                raise ValueError(
                    f"Example {i} must have either 'query' or 'question' key. "
                    f"Found keys: {list(example.keys())}"
                )

            code = example.get('code')
            if not code:
                logger.error(f"Example {i} missing 'code' key: {example.keys()}")
                raise ValueError(f"Example {i} must have 'code' key")

            documents.append(query)
            ids.append(f"example_{i}_{abs(hash(query))}")
            metadatas.append({
                'query': query,  # Store as 'query' in new collection
                'code': code,
                'answer': example.get('answer', ''),
                'description': example.get('description', ''),
                'tags': ','.join(example.get('tags', []))
            })

        try:
            self.examples_collection.add(
                documents=documents,
                ids=ids,
                metadatas=metadatas
            )
            logger.info(f"Successfully added {len(examples)} examples to examples_collection")
        except Exception as e:
            # Handle duplicates by deleting and re-adding
            logger.warning(f"Failed to add examples, attempting to handle duplicates: {e}")
            try:
                self.examples_collection.delete(ids=ids)
                self.examples_collection.add(
                    documents=documents,
                    ids=ids,
                    metadatas=metadatas
                )
                logger.info(f"Successfully added {len(examples)} examples after handling duplicates")
            except Exception as e2:
                logger.error(f"Failed to add example queries even after duplicate handling: {e2}")
                raise RuntimeError(f"Failed to add example queries to vector store: {e2}") from e2

    @retry_on_failure(max_retries=3, delay=0.5)
    def add_business_terms(self, glossary: Dict[str, str]) -> None:
        """
        Add business glossary to vector store.

        Args:
            glossary: Dictionary mapping business terms to definitions

        Raises:
            ValueError: If glossary is invalid
            RuntimeError: If ChromaDB operation fails after retries
        """
        if not self.glossary_collection:
            logger.warning("glossary_collection not initialized, skipping add_business_terms")
            return

        if not glossary:
            logger.warning("add_business_terms called with empty glossary")
            return

        if not isinstance(glossary, dict):
            raise ValueError(f"glossary must be a dict, got {type(glossary)}")

        documents = []
        ids = []
        metadatas = []

        for term, definition in glossary.items():
            if not isinstance(term, str) or not isinstance(definition, str):
                logger.warning(f"Skipping invalid glossary entry: {term} = {definition}")
                continue

            if not term.strip() or not definition.strip():
                logger.warning(f"Skipping empty glossary entry: '{term}' = '{definition}'")
                continue

            documents.append(f"{term}: {definition}")
            ids.append(f"term_{term.lower().replace(' ', '_')}")
            metadatas.append({'term': term, 'definition': definition})

        if not documents:
            logger.error("No valid glossary entries found")
            raise ValueError("No valid glossary entries found")

        try:
            self.glossary_collection.add(
                documents=documents,
                ids=ids,
                metadatas=metadatas
            )
            logger.info(f"Successfully added {len(documents)} business terms to glossary")
        except Exception as e:
            # Handle duplicates
            logger.warning(f"Failed to add business terms, attempting to handle duplicates: {e}")
            try:
                self.glossary_collection.delete(ids=ids)
                self.glossary_collection.add(
                    documents=documents,
                    ids=ids,
                    metadatas=metadatas
                )
                logger.info(f"Successfully added {len(documents)} business terms after handling duplicates")
            except Exception as e2:
                logger.error(f"Failed to add business terms even after duplicate handling: {e2}")
                raise RuntimeError(f"Failed to add business terms to vector store: {e2}") from e2

    def retrieve_context(self, query: str, top_k: int = 5) -> Dict:
        """
        Retrieve relevant context for a query using semantic search.

        Args:
            query: Natural language query
            top_k: Number of examples to retrieve

        Returns:
            dict: Context including schema, examples, and glossary results

        Raises:
            ValueError: If query is invalid
        """
        if not query or not isinstance(query, str):
            raise ValueError(f"query must be a non-empty string, got: {type(query)}")

        if not query.strip():
            raise ValueError("query cannot be empty or whitespace")

        if top_k < 1:
            raise ValueError(f"top_k must be >= 1, got: {top_k}")

        context = {}

        # Search schema
        if self.schema_collection and self.schema_collection.count() > 0:
            try:
                schema_results = self.schema_collection.query(
                    query_texts=[query],
                    n_results=min(2, self.schema_collection.count())
                )
                context['schema'] = schema_results
                logger.debug(f"Retrieved {len(schema_results.get('documents', [[]])[0])} schema documents")
            except Exception as e:
                logger.warning(f"Failed to retrieve schema context: {e}")
                context['schema'] = None

        # Search examples
        if self.examples_collection and self.examples_collection.count() > 0:
            try:
                example_results = self.examples_collection.query(
                    query_texts=[query],
                    n_results=min(top_k, self.examples_collection.count())
                )
                context['examples'] = example_results
                logger.debug(f"Retrieved {len(example_results.get('documents', [[]])[0])} example documents")
            except Exception as e:
                logger.warning(f"Failed to retrieve example context: {e}")
                context['examples'] = None

        # Search glossary
        if self.glossary_collection and self.glossary_collection.count() > 0:
            try:
                glossary_results = self.glossary_collection.query(
                    query_texts=[query],
                    n_results=min(3, self.glossary_collection.count())
                )
                context['glossary'] = glossary_results
                logger.debug(f"Retrieved {len(glossary_results.get('documents', [[]])[0])} glossary terms")
            except Exception as e:
                logger.warning(f"Failed to retrieve glossary context: {e}")
                context['glossary'] = None

        return context

    @retry_on_failure(max_retries=2, delay=0.3)
    def cache_query(self, query: str, code: str, result: str, success: bool) -> None:
        """
        Cache successful query for future use.

        Args:
            query: Natural language query
            code: Generated Python code
            result: Execution result
            success: Whether execution was successful

        Raises:
            ValueError: If inputs are invalid
        """
        if not self.cache_collection:
            logger.debug("cache_collection not initialized, skipping cache_query")
            return

        if not success:
            logger.debug("Query execution was not successful, skipping cache")
            return

        if not query or not isinstance(query, str):
            raise ValueError(f"query must be a non-empty string, got: {type(query)}")

        if not code or not isinstance(code, str):
            raise ValueError(f"code must be a non-empty string, got: {type(code)}")

        try:
            cache_id = f"cache_{abs(hash(query))}"
            self.cache_collection.add(
                documents=[query],
                ids=[cache_id],
                metadatas=[{
                    'code': code,
                    'result': str(result),
                    'timestamp': str(pd.Timestamp.now())
                }]
            )
            logger.debug(f"Successfully cached query: {query[:50]}...")
        except Exception as e:
            # Handle duplicates
            logger.debug(f"Cache entry exists, updating: {e}")
            try:
                cache_id = f"cache_{abs(hash(query))}"
                self.cache_collection.delete(ids=[cache_id])
                self.cache_collection.add(
                    documents=[query],
                    ids=[cache_id],
                    metadatas=[{
                        'code': code,
                        'result': str(result),
                        'timestamp': str(pd.Timestamp.now())
                    }]
                )
                logger.debug(f"Successfully updated cached query")
            except Exception as e2:
                logger.warning(f"Failed to cache query: {e2}")

    def search_cache(self, query: str, threshold: float = 0.9) -> Optional[Dict]:
        """
        Search cache for similar queries (semantic cache).

        Args:
            query: Natural language query
            threshold: Similarity threshold (0-1)

        Returns:
            dict: Cached code and result if found, None otherwise

        Raises:
            ValueError: If inputs are invalid
        """
        if not query or not isinstance(query, str):
            raise ValueError(f"query must be a non-empty string, got: {type(query)}")

        if not 0 <= threshold <= 1:
            raise ValueError(f"threshold must be between 0 and 1, got: {threshold}")

        if not self.cache_collection:
            logger.debug("cache_collection not initialized, cache miss")
            return None

        if self.cache_collection.count() == 0:
            logger.debug("cache_collection is empty, cache miss")
            return None

        try:
            results = self.cache_collection.query(
                query_texts=[query],
                n_results=1
            )

            if results['distances'] and results['distances'][0]:
                # ChromaDB returns distance, lower is better
                # Convert to similarity score
                distance = results['distances'][0][0]
                similarity = 1 - distance

                logger.debug(f"Cache similarity: {similarity:.3f} (threshold: {threshold})")

                if similarity >= threshold:
                    logger.info(f"Cache HIT! Similarity: {similarity:.3f}")
                    return {
                        'code': results['metadatas'][0][0]['code'],
                        'result': results['metadatas'][0][0]['result'],
                        'cached': True,
                        'similarity': similarity
                    }
                else:
                    logger.debug(f"Cache MISS: similarity {similarity:.3f} below threshold {threshold}")
        except Exception as e:
            logger.warning(f"Failed to search cache: {e}")

        return None


def initialize_vector_store(force_reload: bool = False) -> VectorStore:
    """
    Initialize and populate vector store with default examples.

    Args:
        force_reload: If True, clear existing data and reload

    Returns:
        VectorStore: Initialized vector store
    """
    from .example_queries import ExampleQueryLibrary

    vector_store = VectorStore()

    # Check if already populated
    if vector_store.count() > 0 and not force_reload:
        print(
            f"Vector store already initialized with {vector_store.count()} examples"
        )
        return vector_store

    # Clear and reload if forced
    if force_reload:
        vector_store.clear()

    # Get examples from library
    examples = ExampleQueryLibrary.get_examples()

    # Format for vector store
    formatted_examples = []
    for ex in examples:
        formatted_examples.append({
            'question': ex['query'],
            'code': ex['code'],
            'description': ex.get('description', ''),
            'tags': ex.get('tags', [])
        })

    vector_store.add_examples(formatted_examples)

    print(
        f"Vector store initialized with {vector_store.count()} examples"
    )

    return vector_store
