import os
from typing import List, Dict, Any, Optional
from dotenv import load_dotenv
from utils.logger import app_logger
from utils.gemini_llm import GeminiLLM

load_dotenv()

class QueryExpander:
    """
    Expand user queries to improve retrieval performance
    """
    
    def __init__(self):
        """Initialize the query expander with existing Gemini LLM implementation"""
        # Use the existing GeminiLLM implementation
        self.llm = GeminiLLM(model_name="gemini-pro", temperature=0.2)
        
        # Company context for query expansion
        self.company_context = """
        MangoIT Solutions is a web and mobile app development company that offers various services including:
        - Web Development (PHP, Laravel, CodeIgniter, WordPress, Magento)
        - Mobile App Development (iOS, Android, React Native)
        - eCommerce Development (Magento, Shopify, WooCommerce)
        - Custom Software Development
        - AI/ML Development
        - Digital Marketing
        - UI/UX Design
        
        They work with technologies like PHP, Python, JavaScript, React, Angular, Node.js, and various frameworks.
        """
    
    async def expand_query(self, query: str) -> str:
        """
        Expand a user query to improve retrieval performance
        
        Args:
            query: Original user query
            
        Returns:
            Expanded query
        """
        try:
            prompt = f"""
            I need to expand this user query to improve retrieval from a vector database.
            
            Original query: "{query}"
            
            Company context:
            {self.company_context}
            
            Please rewrite the query to be more specific and include relevant keywords that might help in retrieving better results.
            If the query is about technologies, include specific technology names.
            If the query is about services, include specific service names.
            
            The expanded query should be comprehensive but natural sounding.
            
            Return ONLY the expanded query text without any explanations or additional text.
            """
            
            # Use the existing GeminiLLM implementation
            response = await self.llm.agenerate(prompt)
            expanded_query = response.strip()
            
            app_logger.info(f"Original query: {query}")
            app_logger.info(f"Expanded query: {expanded_query}")
            
            return expanded_query
        except Exception as e:
            app_logger.error(f"Error expanding query: {str(e)}")
            # Fall back to original query if expansion fails
            return query
    
    async def expand_query_with_metadata(self, query: str, metadata_context: Optional[Dict[str, Any]] = None) -> str:
        """
        Expand a user query with additional metadata context
        
        Args:
            query: Original user query
            metadata_context: Additional metadata context to use for expansion
            
        Returns:
            Expanded query
        """
        try:
            metadata_str = ""
            if metadata_context:
                metadata_str = "Additional context:\n"
                for key, value in metadata_context.items():
                    if isinstance(value, list):
                        metadata_str += f"- {key}: {', '.join(value)}\n"
                    else:
                        metadata_str += f"- {key}: {value}\n"
            
            prompt = f"""
            I need to expand this user query to improve retrieval from a vector database.
            
            Original query: "{query}"
            
            Company context:
            {self.company_context}
            
            {metadata_str}
            
            Please rewrite the query to be more specific and include relevant keywords that might help in retrieving better results.
            If the query is about technologies, include specific technology names.
            If the query is about services, include specific service names.
            
            The expanded query should be comprehensive but natural sounding.
            
            Return ONLY the expanded query text without any explanations or additional text.
            """
            
            # Use the existing GeminiLLM implementation
            response = await self.llm.agenerate(prompt)
            expanded_query = response.strip()
            
            app_logger.info(f"Original query: {query}")
            app_logger.info(f"Expanded query: {expanded_query}")
            
            return expanded_query
        except Exception as e:
            app_logger.error(f"Error expanding query with metadata: {str(e)}")
            # Fall back to original query if expansion fails
            return query