import gradio as gr
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.chrome.service import Service
from webdriver_manager.chrome import ChromeDriverManager
import google.generativeai as genai
import time
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Configure Gemini
genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
model = genai.GenerativeModel('gemini-2.0-flash')

# Initialize Selenium WebDriver
def init_driver():
    chrome_options = Options()
    chrome_options.add_argument("--headless")
    chrome_options.add_argument("--disable-gpu")
    chrome_options.add_argument("--window-size=1920x1080")
    chrome_options.add_argument("--no-sandbox")
    chrome_options.add_argument("--disable-dev-shm-usage")
    service = Service(ChromeDriverManager().install())
    driver = webdriver.Chrome(service=service, options=chrome_options)
    return driver

# Enhanced web search with Selenium
def search_with_selenium(query):
    driver = init_driver()
    try:
        print(f"Searching for: {query}")
        driver.get(f"https://www.google.com/search?q={query}")
        time.sleep(2)  # Wait for results to load
        
        # Accept cookies if popup appears (for EU users)
        try:
            driver.find_element(By.ID, "L2AGLb").click()
            time.sleep(1)
        except:
            pass
        
        # Extract search results
        results = []
        search_items = driver.find_elements(By.CSS_SELECTOR, 'div.g')
        
        for item in search_items[:5]:  # Get top 5 results
            try:
                title = item.find_element(By.CSS_SELECTOR, 'h3').text
                link = item.find_element(By.CSS_SELECTOR, 'a').get_attribute('href')
                
                # Try different selectors for snippet
                snippet = ""
                try:
                    snippet = item.find_element(By.CSS_SELECTOR, 'div.IsZvec').text
                except:
                    try:
                        snippet = item.find_element(By.CSS_SELECTOR, '.VwiC3b').text
                    except:
                        pass
                
                if title and link:
                    results.append({
                        'title': title,
                        'link': link,
                        'snippet': snippet[:200] + "..." if snippet else "No description available"
                    })
            except Exception as e:
                print(f"Error extracting result: {e}")
                continue
                
        return results
    except Exception as e:
        print(f"Search error: {e}")
        return []
    finally:
        driver.quit()

# Chatbot function with enhanced error handling
def chatbot(message, history, use_web_search):
    try:
        search_results = []
        context = message
        
        # Perform web search if enabled
        if use_web_search:
            with gr.Blocks(analytics_enabled=False):
                with gr.Row():
                    gr.Markdown("🔍 Searching the web...")
            
            search_results = search_with_selenium(message)
            if search_results:
                context += "\n\nWeb search results:\n"
                for i, result in enumerate(search_results, 1):
                    context += f"{i}. {result['title']}\n{result['snippet']}\n\n"
        
        # Generate response
        response = model.generate_content(context)
        bot_message = response.text
        
        # Format response with search results
        if search_results:
            bot_message += "\n\n📚 Sources:\n"
            for result in search_results:
                bot_message += f"- [{result['title']}]({result['link']})\n"
        
        return bot_message
    except Exception as e:
        return f"⚠️ Error: {str(e)}"

# Gradio interface with enhanced UI
with gr.Blocks(
    title="Gemini AI Chatbot with Web Search",
    theme=gr.themes.Soft(),
    css=".gradio-container {max-width: 800px !important}"
) as demo:
    
    gr.Markdown("""
    # 💎 Gemini AI Chatbot
    **Powered by Google Gemini Pro + Real-time Web Search**
    """)
    
    with gr.Row():
        with gr.Column(scale=4):
            chatbot = gr.Chatbot(
                height=500,
                bubble_full_width=False,
                avatar_images=(
                    "user.png", 
                    "bot.png"
                )
            )
            msg = gr.Textbox(
                label="Your Message",
                placeholder="Ask me anything...",
                lines=2
            )
            
            with gr.Row():
                submit_btn = gr.Button("Send", variant="primary")
                clear_btn = gr.Button("Clear")
                web_search = gr.Checkbox(
                    label="Enable Web Search",
                    value=True,
                    interactive=True
                )
        
        with gr.Column(scale=1):
            gr.Markdown("""
            ### 🌐 Web Search Features
            - Real-time Google results
            - Source citations
            - Toggle on/off
            
            ### ⚙️ Controls
            - Press Enter or click Send
            - Clear resets conversation
            """)
    
    # Event handlers
    def respond(message, chat_history, use_search):
        bot_message = chatbot(message, chat_history, use_search)
        chat_history.append((message, bot_message))
        return "", chat_history
    
    msg.submit(respond, [msg, chatbot, web_search], [msg, chatbot])
    submit_btn.click(respond, [msg, chatbot, web_search], [msg, chatbot])
    clear_btn.click(lambda: None, None, chatbot, queue=False)

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        favicon_path="favicon.ico"
    )