"""
NLQ System Initialization Helper.

Provides utilities for one-time initialization of the Natural Language Query system.
Populates vector database with metadata, examples, and business glossary.
"""

from typing import Tuple, Optional
from .semantic_layer import SemanticLayer
from .example_queries import ExampleQueryLibrary
from .vector_store import VectorStore


def initialize_nlq_system(
    eventstream=None,
    persist_directory: Optional[str] = None,
    force_reload: bool = False,
    verbose: bool = True
) -> Tuple[VectorStore, SemanticLayer]:
    """
    One-time initialization of NLQ system.
    Populates vector database with metadata.

    Args:
        eventstream: Optional Eventstream object for schema extraction
        persist_directory: Optional directory for ChromaDB persistence
        force_reload: If True, clear existing data and reload
        verbose: Print progress messages

    Returns:
        tuple: (VectorStore, SemanticLayer) instances

    Example:
        >>> from outhad_edge import datasets
        >>> stream = datasets.load_simple_shop()
        >>> vector_store, semantic_layer = initialize_nlq_system(stream)
    """

    if verbose:
        print("=" * 70)
        print("Initializing NLQ system...")
        print("=" * 70)

    # Create semantic layer
    if verbose:
        print("\n[1/5] Creating semantic layer...")

    if eventstream:
        semantic_layer = SemanticLayer(eventstream)
        if verbose:
            print(f"✓ Semantic layer created for eventstream")
    else:
        # Create empty semantic layer (will need eventstream later)
        semantic_layer = None
        if verbose:
            print("⚠ No eventstream provided, skipping schema metadata")

    # Create vector store
    if verbose:
        print("\n[2/5] Creating vector store...")

    vector_store = VectorStore(persist_directory=persist_directory)

    if force_reload:
        if verbose:
            print("⚠ Force reload enabled, clearing existing data...")
        vector_store.clear()

    if verbose:
        print(f"✓ Vector store created at: {vector_store.persist_directory}")

    # Add schema metadata (if eventstream provided)
    if semantic_layer:
        if verbose:
            print("\n[3/5] Adding schema metadata...")

        try:
            vector_store.add_schema_metadata(semantic_layer)
            if verbose:
                schema_count = vector_store.schema_collection.count() if vector_store.schema_collection else 0
                print(f"✓ Added schema metadata ({schema_count} documents)")
        except Exception as e:
            if verbose:
                print(f"✗ Failed to add schema metadata: {e}")
    else:
        if verbose:
            print("\n[3/5] Skipping schema metadata (no eventstream)")

    # Add example queries
    if verbose:
        print("\n[4/5] Adding example queries...")

    try:
        examples = ExampleQueryLibrary.get_examples()

        # Add to both collections for backward compatibility
        vector_store.add_examples(examples)  # Main collection (backward compat)

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

        vector_store.add_example_queries(formatted_examples)  # New collection

        if verbose:
            print(f"✓ Added {len(examples)} example queries")
            if vector_store.examples_collection:
                print(f"  - Examples collection: {vector_store.examples_collection.count()} documents")
            print(f"  - Main collection: {vector_store.collection.count()} documents")

    except Exception as e:
        if verbose:
            print(f"✗ Failed to add examples: {e}")

    # Add business glossary
    if semantic_layer:
        if verbose:
            print("\n[5/5] Adding business glossary...")

        try:
            vector_store.add_business_terms(semantic_layer.business_glossary)
            if verbose:
                glossary_count = vector_store.glossary_collection.count() if vector_store.glossary_collection else 0
                print(f"✓ Added business glossary ({glossary_count} terms)")
        except Exception as e:
            if verbose:
                print(f"✗ Failed to add glossary: {e}")
    else:
        if verbose:
            print("\n[5/5] Skipping business glossary (no eventstream)")

    if verbose:
        print("\n" + "=" * 70)
        print("✓ NLQ system initialized successfully!")
        print("=" * 70)

        print("\nCollection Status:")
        print(f"  - Main collection: {vector_store.collection.count()} documents")
        if vector_store.schema_collection:
            print(f"  - Schema metadata: {vector_store.schema_collection.count()} documents")
        if vector_store.examples_collection:
            print(f"  - Example queries: {vector_store.examples_collection.count()} documents")
        if vector_store.glossary_collection:
            print(f"  - Business glossary: {vector_store.glossary_collection.count()} terms")
        if vector_store.cache_collection:
            print(f"  - Query cache: {vector_store.cache_collection.count()} cached queries")

        print("\nUsage:")
        print("  >>> from outhad_edge import datasets")
        print("  >>> stream = datasets.load_simple_shop()")
        print("  >>> nlq = stream.nlq(model='gpt-4')")
        print("  >>> result = nlq.ask('What is the conversion rate?')")

    return vector_store, semantic_layer


def reinitialize_vector_store(
    eventstream=None,
    persist_directory: Optional[str] = None,
    verbose: bool = True
) -> Tuple[VectorStore, Optional[SemanticLayer]]:
    """
    Reinitialize vector store by clearing and reloading all data.

    Args:
        eventstream: Optional Eventstream object
        persist_directory: Optional directory for ChromaDB persistence
        verbose: Print progress messages

    Returns:
        tuple: (VectorStore, SemanticLayer) instances
    """
    return initialize_nlq_system(
        eventstream=eventstream,
        persist_directory=persist_directory,
        force_reload=True,
        verbose=verbose
    )


def check_nlq_status(vector_store: VectorStore) -> dict:
    """
    Check the current status of NLQ system initialization.

    Args:
        vector_store: VectorStore instance to check

    Returns:
        dict: Status information about each collection
    """
    status = {
        'initialized': False,
        'collections': {}
    }

    # Check main collection
    status['collections']['main'] = {
        'exists': vector_store.collection is not None,
        'count': vector_store.collection.count() if vector_store.collection else 0
    }

    # Check schema collection
    status['collections']['schema'] = {
        'exists': vector_store.schema_collection is not None,
        'count': vector_store.schema_collection.count() if vector_store.schema_collection else 0
    }

    # Check examples collection
    status['collections']['examples'] = {
        'exists': vector_store.examples_collection is not None,
        'count': vector_store.examples_collection.count() if vector_store.examples_collection else 0
    }

    # Check glossary collection
    status['collections']['glossary'] = {
        'exists': vector_store.glossary_collection is not None,
        'count': vector_store.glossary_collection.count() if vector_store.glossary_collection else 0
    }

    # Check cache collection
    status['collections']['cache'] = {
        'exists': vector_store.cache_collection is not None,
        'count': vector_store.cache_collection.count() if vector_store.cache_collection else 0
    }

    # Determine if initialized
    examples_count = status['collections']['examples']['count']
    status['initialized'] = examples_count > 0

    return status
