import gradio as gr
import os
from google import genai
from google.genai.types import Tool, GenerateContentConfig, GoogleSearch, UrlContext

# Initialize the Gemini client
def initialize_client(api_key):
    """Initialize the Gemini client with API key"""
    try:
        client = genai.Client(api_key=api_key)
        return client, None
    except Exception as e:
        return None, str(e)

def search_with_context(api_key, query, url=None, enable_google_search=True, enable_url_context=True):
    """
    Perform search using Gemini with optional URL context and Google search
    
    Args:
        api_key: Google Gemini API key
        query: Search query or question
        url: Optional URL for context
        enable_google_search: Whether to enable Google search
        enable_url_context: Whether to enable URL context
    """
    
    if not api_key:
        return "❌ Please provide a valid API key", "", ""
    
    if not query:
        return "❌ Please provide a search query", "", ""
    
    try:
        # Initialize client
        client, error = initialize_client(api_key)
        if error:
            return f"❌ Error initializing client: {error}", "", ""
        
        model_id = "gemini-2.5-flash"
        
        # Configure tools based on user selection
        tools = []
        if enable_url_context:
            tools.append({"url_context": {}})
        if enable_google_search:
            tools.append({"google_search": {}})
        
        if not tools:
            return "❌ Please enable at least one tool (URL Context or Google Search)", "", ""
        
        # Modify query to include URL if provided
        final_query = query
        if url and enable_url_context:
            final_query = f"{query} : {url}"
        
        # Generate content
        response = client.models.generate_content(
            model=model_id,
            contents=final_query,
            config=GenerateContentConfig(
                tools=tools,
            )
        )
        
        # Extract response text
        response_text = ""
        for part in response.candidates[0].content.parts:
            if hasattr(part, 'text') and part.text:
                response_text += part.text + "\n"
        
        # Extract URL context metadata
        url_metadata = ""
        try:
            if hasattr(response.candidates[0], 'url_context_metadata') and response.candidates[0].url_context_metadata:
                url_metadata = "📋 **URL Context Metadata:**\n\n"
                for metadata in response.candidates[0].url_context_metadata:
                    url_metadata += f"• **URL:** {metadata.url}\n"
                    url_metadata += f"• **Title:** {metadata.title}\n"
                    if hasattr(metadata, 'snippet'):
                        url_metadata += f"• **Snippet:** {metadata.snippet}\n"
                    url_metadata += "\n"
        except Exception as e:
            url_metadata = f"URL metadata extraction error: {str(e)}"
        
        # Extract search metadata (if available)
        search_metadata = ""
        try:
            if hasattr(response.candidates[0], 'grounding_metadata') and response.candidates[0].grounding_metadata:
                search_metadata = "🔍 **Search Metadata:**\n\n"
                # Add search metadata details here if available
                search_metadata += str(response.candidates[0].grounding_metadata)
        except Exception as e:
            search_metadata = f"Search metadata extraction error: {str(e)}"
        
        if not response_text.strip():
            response_text = "No response generated. Please check your query and try again."
        
        return response_text.strip(), url_metadata, search_metadata
        
    except Exception as e:
        return f"❌ Error: {str(e)}", "", ""

# Create the Gradio interface
def create_gradio_interface():
    """Create and configure the Gradio interface"""
    
    with gr.Blocks(
        title="🔍 Gemini URL Context & Search App",
        theme=gr.themes.Soft(),
        css="""
        .main-container { max-width: 1200px; margin: 0 auto; }
        .response-box { min-height: 200px; }
        .metadata-box { min-height: 150px; }
        """
    ) as app:
        
        gr.HTML("""
        <div style='text-align: center; margin-bottom: 20px;'>
            <h1>🔍 Gemini URL Context & Search App</h1>
            <p>Search and analyze content using Google's Gemini AI with URL context and web search capabilities</p>
        </div>
        """)
        
        with gr.Row():
            with gr.Column(scale=2):
                # Input section
                with gr.Group():
                    gr.HTML("<h3>🔑 Configuration</h3>")
                    api_key_input = gr.Textbox(
                        label="Gemini API Key",
                        placeholder="Enter your Google Gemini API key",
                        type="password",
                        info="Get your API key from Google AI Studio"
                    )
                
                with gr.Group():
                    gr.HTML("<h3>📝 Query & Context</h3>")
                    query_input = gr.Textbox(
                        label="Search Query",
                        placeholder="Ask a question or enter your search query...",
                        lines=3,
                        info="Your question or search query"
                    )
                    
                    url_input = gr.Textbox(
                        label="URL for Context (Optional)",
                        placeholder="https://example.com/article",
                        info="Provide a URL to analyze or use as context for your query"
                    )
                
                with gr.Group():
                    gr.HTML("<h3>⚙️ Tool Settings</h3>")
                    with gr.Row():
                        enable_url_context = gr.Checkbox(
                            label="Enable URL Context",
                            value=True,
                            info="Use URL content as context"
                        )
                        enable_google_search = gr.Checkbox(
                            label="Enable Google Search",
                            value=True,
                            info="Use Google search for additional context"
                        )
                
                search_button = gr.Button(
                    "🚀 Search & Analyze",
                    variant="primary",
                    size="lg"
                )
            
            with gr.Column(scale=3):
                # Output section
                with gr.Group():
                    gr.HTML("<h3>💬 AI Response</h3>")
                    response_output = gr.Textbox(
                        label="Generated Response",
                        lines=15,
                        elem_classes=["response-box"],
                        show_copy_button=True
                    )
                
                with gr.Accordion("📋 URL Context Details", open=False):
                    url_metadata_output = gr.Textbox(
                        label="URL Context Metadata",
                        lines=8,
                        elem_classes=["metadata-box"],
                        show_copy_button=True
                    )
                
                with gr.Accordion("🔍 Search Details", open=False):
                    search_metadata_output = gr.Textbox(
                        label="Search Metadata",
                        lines=8,
                        elem_classes=["metadata-box"],
                        show_copy_button=True
                    )
        
        # Examples section
        with gr.Accordion("💡 Example Queries", open=False):
            gr.Examples(
                examples=[
                    [
                        "What are the main points discussed in this article?",
                        "https://www.africanhistoryextra.com/p/state-society-and-ethnicity-in-19th",
                        True,
                        True
                    ],
                    [
                        "Summarize the latest developments in AI technology",
                        "",
                        False,
                        True
                    ],
                    [
                        "What are the key features and pricing of this product?",
                        "https://example.com/product-page",
                        True,
                        False
                    ],
                    [
                        "Compare the content of this page with current industry trends",
                        "https://example.com/industry-report",
                        True,
                        True
                    ]
                ],
                inputs=[query_input, url_input, enable_url_context, enable_google_search],
                label="Click on any example to load it"
            )
        
        # Event handler
        search_button.click(
            fn=search_with_context,
            inputs=[
                api_key_input,
                query_input,
                url_input,
                enable_google_search,
                enable_url_context
            ],
            outputs=[
                response_output,
                url_metadata_output,
                search_metadata_output
            ]
        )
        
        # Add footer
        gr.HTML("""
        <div style='text-align: center; margin-top: 30px; padding: 20px; border-top: 1px solid #ddd;'>
            <p style='color: #666;'>
                🚀 Built with Gradio and Google Gemini API<br>
                💡 Tip: Use URL context for analyzing specific web pages, and Google search for broader queries
            </p>
        </div>
        """)
    
    return app

# Main execution
if __name__ == "__main__":
    # Create and launch the app
    app = create_gradio_interface()
    
    app.launch(
        debug=True,
        mcp_server=True
    )