#!/usr/bin/env python3
"""
CLI for Knowledge Base operations.
"""

import argparse
import json
import logging
import os
import sys
from pathlib import Path
from typing import Optional

# Import txtai
from txtai.app import Application
from txtai.pipeline import Extractor

# Import settings
from .settings import Settings

# Configure logging
logger = logging.getLogger(__name__)

def setup_logging(debug: bool = False):
    """Set up logging configuration.
    
    Args:
        debug: Whether to enable debug logging
    """
    level = logging.DEBUG if debug else logging.INFO
    logging.basicConfig(
        level=level,
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
        handlers=[logging.StreamHandler()]
    )

def find_config_file() -> Optional[str]:
    """
    Find a configuration file in standard locations.
    
    Returns:
        Path to the configuration file if found, None otherwise.
    """
    # Check environment variable first
    if os.environ.get("KB_CONFIG"):
        config_path = os.environ.get("KB_CONFIG")
        if os.path.exists(config_path):
            return config_path
    
    # Check standard locations
    search_paths = [
        "./config.yaml",
        "./config.yml",
        Path.home() / ".config" / "knowledge-base" / "config.yaml",
        Path.home() / ".config" / "knowledge-base" / "config.yml",
        Path.home() / ".knowledge-base" / "config.yaml",
        Path.home() / ".knowledge-base" / "config.yml",
    ]
    
    for path in search_paths:
        if os.path.exists(path):
            return str(path)
    
    return None

def create_application(config_path: Optional[str] = None) -> Application:
    """
    Create a txtai application with the specified configuration.
    
    Args:
        config_path: Path to configuration file
        
    Returns:
        txtai.app.Application: Application instance
    """
    # Use provided config if available
    if config_path:
        # Convert to absolute path if it's a relative path
        if not os.path.isabs(config_path):
            config_path = os.path.abspath(config_path)
        
        if os.path.exists(config_path):
            logger.info(f"Loading configuration from {config_path}")
            try:
                # Create application directly from YAML file path
                app = Application(config_path)
                
                # Log configuration details
                if hasattr(app.embeddings, 'graph') and app.embeddings.graph:
                    logger.info("Graph configuration found in embeddings")
                
                # Log index path
                if hasattr(app, 'config') and 'path' in app.config:
                    logger.info(f"Index will be stored at: {app.config['path']}")
                
                return app
            except Exception as e:
                logger.error(f"Error loading configuration from {config_path}: {e}")
                logger.warning("Falling back to default configuration")
        else:
            logger.warning(f"Configuration file not found: {config_path}")
            logger.warning("Falling back to default configuration")
    else:
        logger.info("No configuration file specified, using default configuration")
    
    # If no config provided or loading failed, use default settings
    logger.info("Creating application with default configuration")
    
    # Get settings
    settings = Settings(config_path)
    
    # Create default configuration
    config = {
        "path": ".txtai/index",  # Default index path
        "writable": True,  # Enable index writing
        "content": True,   # Store document content
        "embeddings": {
            "path": settings.get("model_path", "sentence-transformers/all-MiniLM-L6-v2"),
            "gpu": settings.get("model_gpu", True),
            "normalize": settings.get("model_normalize", True),
            "content": True,  # Store document content
            "writable": True   # Enable index writing
        },
        "search": {
            "hybrid": settings.get("hybrid_search", False)
        }
    }
    
    return Application(config)

def build_command(args):
    """
    Handle build command.
    
    Args:
        args: Command-line arguments
        
    The build command processes input documents and builds a searchable index.
    When used with --update, it will add to or update an existing index instead of rebuilding it.
    It can also export the built index to a compressed tar.gz file for portability.
    """
    # Check if required arguments are provided
    if not ((hasattr(args, 'input') and args.input) or 
            (hasattr(args, 'json_input') and args.json_input)):
        logger.error("Error: No input sources provided.")
        logger.error("Please provide at least one input source using --input or --json_input")
        print("\nBuild command usage:")
        print("  python -m kb_builder build --input PATH [PATH ...] [--config CONFIG] [--export EXPORT_PATH] [--update]")
        print("  python -m kb_builder build --json_input JSON_FILE [--config CONFIG] [--export EXPORT_PATH] [--update]")
        print("\nOptions:")
        print("  --input PATH       Path to input files or directories")
        print("  --json_input PATH  Path to JSON file containing a list of documents")
        print("  --extensions EXT   Comma-separated list of file extensions to include")
        print("  --config PATH      Path to configuration file")
        print("  --export PATH      Export the built index to a compressed tar.gz file")
        print("  --update           Update existing index instead of rebuilding it")
        return
        
    # Use config from args or try to find a default config
    config_path = args.config if hasattr(args, 'config') and args.config else None
    
    # If no config provided via args, try to find one
    if not config_path:
        config_path = find_config_file()
    
    if config_path:
        # Convert to absolute path if it's a relative path
        if not os.path.isabs(config_path):
            config_path = os.path.abspath(config_path)
            
        if not os.path.exists(config_path):
            logger.error(f"Configuration file not found: {config_path}")
            logger.error("Please provide a valid path to a configuration file")
            return
        
        logger.info(f"Using configuration from {config_path}")
    else:
        logger.warning("No configuration file specified, using default settings")
    
    # Create application
    app = create_application(config_path)
    
    # Verify textractor pipeline exists
    if "textractor" not in app.pipelines:
        logger.error("No textractor pipeline configured in YAML. Please add a 'textractor' section to your configuration.")
        logger.error("Example: textractor:\n  paragraphs: true\n  minlength: 100")
        return
    
    # Process documents
    documents = []
    
    # Process JSON input if provided
    if args.json_input:
        try:
            with open(args.json_input, 'r') as f:
                json_data = json.load(f)
                
            # Check if it's a list of documents
            if isinstance(json_data, list):
                documents.extend(json_data)
                logger.info(f"Loaded {len(json_data)} documents from {args.json_input}")
            else:
                logger.error(f"Invalid JSON format in {args.json_input}. Expected a list of documents.")
        except Exception as e:
            logger.error(f"Error loading JSON from {args.json_input}: {e}")
    
    # Process file/directory inputs
    if args.input:
        # Parse extensions
        extensions = None
        if args.extensions:
            # Convert comma-separated string to set of extensions
            extensions = set(ext.strip().lower() for ext in args.extensions.split(","))
            # Add leading dot if not present
            extensions = {ext if ext.startswith('.') else f'.{ext}' for ext in extensions}
        
        for input_path in args.input:
            path = Path(input_path)
            
            if path.is_file():
                logger.info(f"Processing file: {path}")
                try:
                    # Extract text using textractor pipeline
                    segments = app.pipelines["textractor"](str(path))
                    
                    # Create documents with metadata
                    for i, text in enumerate(segments):
                        doc_id = f"{path.stem}_{i}"
                        documents.append({
                            "id": doc_id,
                            "text": text,
                            "metadata": {
                                "source": str(path),
                                "index": i,
                                "total": len(segments)
                            }
                        })
                    logger.info(f"Extracted {len(segments)} segments from {path}")
                except Exception as e:
                    logger.error(f"Error processing file {path}: {e}")
            
            elif path.is_dir():
                logger.info(f"Processing directory: {path}")
                try:
                    # Find all files in directory
                    files = []
                    if extensions:
                        for ext in extensions:
                            files.extend(path.glob(f"**/*{ext}"))
                    else:
                        files = list(path.glob("**/*"))
                    
                    # Filter out directories
                    files = [f for f in files if f.is_file()]
                    
                    logger.info(f"Found {len(files)} files in directory {path}")
                    
                    # Process each file
                    for file_path in files:
                        try:
                            # Extract text using textractor pipeline
                            segments = app.pipelines["textractor"](str(file_path))
                            
                            # Create documents with metadata
                            for i, text in enumerate(segments):
                                doc_id = f"{file_path.stem}_{i}"
                                documents.append({
                                    "id": doc_id,
                                    "text": text,
                                    "metadata": {
                                        "source": str(file_path),
                                        "index": i,
                                        "total": len(segments)
                                    }
                                })
                        except Exception as e:
                            logger.error(f"Error processing file {file_path}: {e}")
                except Exception as e:
                    logger.error(f"Error processing directory {path}: {e}")
            
            else:
                logger.warning(f"Input path not found: {path}")
    
    # Check if we have documents to process
    if not documents:
        logger.error("No documents found to process")
        return
    
    logger.info(f"Processed {len(documents)} documents")
    
    # Use the application's add method which handles both indexing and saving
    logger.info("Indexing documents...")
    try:
        # Build or update the index
        if hasattr(args, 'update') and args.update:
            # Use upsert to update the existing index
            # The database.stream() method doesn't exist, directly use the documents list
            # Convert metadata to JSON string if it's a dictionary to avoid SQLite binding errors
            documents_for_upsert = []
            for i, doc in enumerate(documents):
                doc_id = doc.get("id", i)
                text = doc.get("text", "")
                metadata = doc.get("metadata")
                # Convert metadata to JSON string if it's a dictionary
                if isinstance(metadata, dict):
                    metadata = json.dumps(metadata)
                documents_for_upsert.append((doc_id, text, metadata))
            
            app.embeddings.upsert(documents_for_upsert)
            logger.info("Documents added to existing index successfully")
        else:
            # Rebuild the entire index
            app.index()
            logger.info("Documents indexed successfully")
        
        # Log if graph was built
        if hasattr(app.embeddings, 'graph') and app.embeddings.graph:
            logger.info("Knowledge graph was automatically built based on YAML configuration")
            
        # Explicitly save the index to the configured path
        if app.config and "path" in app.config:
            save_path = app.config["path"]
            logger.info(f"Saving index to configured path: {save_path}")
            os.makedirs(save_path, exist_ok=True)
            app.embeddings.save(save_path)
            logger.info(f"Index successfully saved to {save_path}")
            
        # Export the index to a tar.gz file if requested
        if hasattr(args, 'export') and args.export:
            export_path = args.export
            # Add .tar.gz extension if not present
            if not export_path.endswith('.tar.gz'):
                export_path = f"{export_path}.tar.gz"
                
            logger.info(f"Exporting index to {export_path}")
            try:
                # Save the embeddings to a tar.gz file
                app.embeddings.save(export_path)
                logger.info(f"Index successfully exported to {export_path}")
                
                # Show how to load this index
                logger.info(f"To load this index, use: app = Application(\"path: {export_path}\")")
            except Exception as e:
                logger.error(f"Error exporting index to {export_path}: {e}")
    except Exception as e:
        logger.error(f"Error indexing documents: {e}")
        return

def retrieve_command(args):
    """
    Handle retrieve command.
    """
    try:
        # Create application
        print(f"Creating application with path: {args.embeddings}")
        app = Application(f"path: {args.embeddings}")

        # Perform search
        print(f"Performing search with query: {args.query}")
        
        # Extract key terms from the query to use for relevance boosting
        query_terms = set(args.query.lower().split())
        # Remove common stop words
        stop_words = {"what", "are", "is", "the", "for", "and", "or", "to", "in", "of", "a", "an"}
        query_terms = query_terms - stop_words
        
        # Perform the search
        results = app.search(args.query, limit=max(10, args.limit * 2), graph=args.graph)  # Get more results initially for filtering

        # Apply generic result enhancement
        if args.graph:
            # For graph results, enhance using centrality and query relevance
            if hasattr(results, 'centrality') and callable(results.centrality):
                # Get all nodes with their centrality scores
                nodes_with_scores = []
                for node_id in results.centrality().keys():
                    node = results.node(node_id)
                    if node and "text" in node:
                        # Base score from centrality
                        score = results.centrality()[node_id]
                        
                        # Boost score based on query term presence
                        text = node["text"].lower()
                        term_matches = sum(1 for term in query_terms if term in text)
                        if term_matches > 0:
                            # Boost proportional to the number of matching terms
                            score *= (1 + (0.2 * term_matches))
                            
                        # Add to candidates if score meets minimum threshold
                        if score >= args.min_similarity:
                            nodes_with_scores.append((node_id, score, node["text"]))
                
                # Sort by enhanced score and limit
                nodes_with_scores.sort(key=lambda x: x[1], reverse=True)
                nodes_with_scores = nodes_with_scores[:args.limit]
                
                # Convert to the format expected by format_graph_results
                graph_results = [{"text": text, "score": score} for _, score, text in nodes_with_scores]
            else:
                # Fallback if centrality not available
                graph_results = []
                for x in list(results)[:args.limit]:
                    if "text" in x:
                        score = x.get("score", 0.5)
                        if score >= args.min_similarity:
                            graph_results.append({"text": x["text"], "score": score})
        else:
            # For regular search results, enhance based on query relevance
            enhanced_results = []
            for result in results:
                if "text" in result and "score" in result:
                    # Base score from search
                    score = result["score"]
                    
                    # Boost score based on query term presence
                    text = result["text"].lower()
                    term_matches = sum(1 for term in query_terms if term in text)
                    if term_matches > 0:
                        # Boost proportional to the number of matching terms
                        score *= (1 + (0.1 * term_matches))
                    
                    # Add to candidates if score meets minimum threshold
                    if score >= args.min_similarity:
                        enhanced_results.append({"text": result["text"], "score": score})
            
            # Sort by enhanced score and limit
            enhanced_results.sort(key=lambda x: x["score"], reverse=True)
            results = enhanced_results[:args.limit]

        # Print results
        if args.graph:
            # Format and print results
            if graph_results:
                formatted_results = format_graph_results(app.embeddings, graph_results, args.query)
                print(formatted_results)
            else:
                print(f"Q:{args.query}")
                print("No results found.\n")
        else:
            print(f"Results for query: '{args.query}'")
            for i, result in enumerate(results):
                print(f"\nResult {i+1}:")
                print(f"  Score: {result['score']:.4f}")
                print(f"  Text: {result['text']}")
                print()

    except Exception as e:
        print(f"Error during retrieval: {e}")
        logger.error(f"Error during retrieval: {e}")



def format_graph_results(embeddings, results, query=None):
    """
    Format graph search results to match graph.ipynb output format.
    """
    output = []
    
    if query:
        output.append(f"Q:{query}")
    
    # Process each result
    for result in results:
        try:
            # Get the text and metadata
            if isinstance(result, dict) and "text" in result:
                text = result["text"]
                # Generate a node id from the content
                words = text.split()[:5]  # Take first 5 words
                node_id = "_".join(w.lower() for w in words if w.isalnum())
            else:
                node = embeddings.graph.node(result)
                if not node:
                    continue
                text = node.get("text", "")
                node_id = result
            
            if not text.strip():
                continue
            
            # Add formatted result
            output.append(f"# {node_id}")
            output.append(text.strip())
            output.append("")  # Empty line between results
            
        except Exception as e:
            logger.error(f"Error processing result {result}: {str(e)}")
            continue
    
    return "\n".join(output)



def main():
    """
    Main entry point for the Knowledge Base CLI.
    """
    parser = argparse.ArgumentParser(description="Knowledge Base CLI")
    
    # Global arguments
    parser.add_argument("--debug", action="store_true", help="Enable debug logging")
    
    subparsers = parser.add_subparsers(title="commands", dest="command", help="Available commands")
    
    # Build command
    build_parser = subparsers.add_parser("build", help="Build embeddings database")
    build_parser.add_argument("--input", type=str, nargs="+", help="Path to input files or directories")
    build_parser.add_argument("--extensions", type=str, help="Comma-separated list of file extensions to include")
    build_parser.add_argument("--json_input", type=str, help="Path to JSON file containing a list of documents")
    build_parser.add_argument("--config", type=str, help="Path to configuration file")
    build_parser.add_argument("--export", type=str, help="Export the built index to a compressed tar.gz file")
    build_parser.add_argument("--update", action="store_true", help="Update existing index instead of rebuilding it")
    build_parser.set_defaults(func=build_command)
    
    # Retrieve command
    retrieve_parser = subparsers.add_parser("retrieve", help="Retrieve information from embeddings database")
    retrieve_parser.add_argument("embeddings", type=str, help="Path to embeddings database")
    retrieve_parser.add_argument("query", type=str, help="Search query")
    retrieve_parser.add_argument("--limit", type=int, default=5, help="Maximum number of results to return")
    retrieve_parser.add_argument("--graph", action="store_true", help="Enable graph search")
    retrieve_parser.add_argument("--min_similarity", type=float, default=0.3, help="Minimum similarity threshold for results")
    retrieve_parser.set_defaults(func=retrieve_command)
    
    # Generate command removed as it's currently unused
    args = parser.parse_args()
    
    # Set up logging
    setup_logging(args.debug if hasattr(args, 'debug') else False)
    
    if args.command:
        args.func(args)
    else:
        parser.print_help()

# Entry point functions have been moved to dedicated scripts in the bin directory

if __name__ == "__main__":
    main()
