"""
Vector Search Tool - Flexible search across one or more product versions.
"""

from typing import Dict, List, Optional
import json


def get_vector_search_tool_definition():
    """Get the OpenAI function definition for the vector search tool."""
    return {
        "type": "function",
        "function": {
            "name": "vector_search",
            "description": "Search for information across one or more product versions. Use this to find specific information or compare across versions.",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "The search query"
                    },
                    "versions": {
                        "type": "array",
                        "items": {
                            "type": "string",
                            "enum": ["harmony_1_2", "harmony_1_5", "harmony_1_6", "harmony_1_8", "chorus_1_1", "general_faq"]
                        },
                        "description": "List of versions to search. Can be a single version or multiple versions for comparison."
                    },
                    "max_results_per_version": {
                        "type": "integer",
                        "description": "Maximum results per version",
                        "default": 5,
                        "minimum": 1,
                        "maximum": 20
                    }
                },
                "required": ["query", "versions"]
            }
        }
    }


def execute_vector_search(vector_store_manager, query: str, versions: List[str], 
                         max_results_per_version: int = 5) -> Dict:
    """
    Execute vector search across one or more versions.
    
    Args:
        vector_store_manager: Instance of VectorStoreManager
        query: Search query
        versions: List of version stores to search
        max_results_per_version: Maximum results per version
    
    Returns:
        Dictionary with search results
    """
    try:
        # If single version, return simple format
        if len(versions) == 1:
            results = vector_store_manager.query_vector_store(
                versions[0], query, max_results_per_version
            )
            
            if not results:
                return {
                    "status": "success",
                    "message": f"No results found in {versions[0]}",
                    "results": []
                }
            
            return {
                "status": "success",
                "message": f"Found {len(results)} results in {versions[0]}",
                "version": versions[0],
                "query": query,
                "results": results
            }
        
        # Multiple versions - return by version
        all_results = {}
        for version in versions:
            results = vector_store_manager.query_vector_store(
                version, query, max_results_per_version
            )
            if results:
                all_results[version] = results
        
        if not all_results:
            return {
                "status": "success",
                "message": f"No results found across versions",
                "results": {}
            }
        
        return {
            "status": "success",
            "message": f"Found results in {len(all_results)} versions",
            "query": query,
            "results": all_results
        }
        
    except Exception as e:
        return {
            "status": "error",
            "message": f"Error searching: {str(e)}",
            "results": {} if len(versions) > 1 else []
        }


def format_search_results_for_context(results: Dict) -> str:
    """Format search results for including in the context."""
    if results["status"] != "success":
        return ""
    
    # Handle single version search
    if "version" in results and results.get("results"):
        formatted = [f"Search results from {results['version']} for '{results['query']}':"]
        
        for i, result in enumerate(results["results"], 1):
            text = result.get("text", "")
            similarity = result.get("similarity", 0)
            formatted.append(f"{i}. {text} (similarity: {similarity:.3f})")
        
        return "\n".join(formatted)
    
    # Handle multi-version search
    if isinstance(results.get("results"), dict):
        formatted = [f"Search results across versions for '{results['query']}':"]
        
        for version, version_results in results["results"].items():
            formatted.append(f"\n{version}:")
            for i, result in enumerate(version_results[:3], 1):
                text = result.get("text", "")
                similarity = result.get("similarity", 0)
                formatted.append(f"  {i}. {text} (similarity: {similarity:.3f})")
        
        return "\n".join(formatted)
    
    return ""


# Keep old function names for compatibility but have them use the new unified tool
def get_multi_version_search_tool_definition():
    """Deprecated - use get_vector_search_tool_definition instead."""
    return get_vector_search_tool_definition()


def execute_multi_version_search(vector_store_manager, query: str, versions: List[str], 
                               max_results_per_version: int = 3) -> Dict:
    """Deprecated - use execute_vector_search instead."""
    return execute_vector_search(vector_store_manager, query, versions, max_results_per_version)