import gradio as gr
from openai import OpenAI
import os
import json

# Initialize OpenAI client with API key and base URL from environment variables
client = OpenAI(
    api_key=os.environ["OPENAI_API_KEY"],
    base_url=os.environ["OPENAI_BASE_URL"]
)

# Define the number of results per page and total results to generate
RESULTS_PER_PAGE = 10
TOTAL_RESULTS = 30  # Generate 30 results to allow pagination

def fetch_search_results(query):
    """Fetch search results from the LLM based on the user's query."""
    if not query.strip():
        return None, "Please enter a search query."
    
    prompt = f"""
    You are a search engine that provides informative and relevant results. For the given query '{query}', 
    generate {TOTAL_RESULTS} search results, each with a title and a snippet that summarizes the information. 
    Format the response as a JSON array of objects, where each object has 'title' and 'snippet' fields.
    Ensure the results are diverse and relevant to the query.
    """
    
    try:
        response = client.chat.completions.create(
            model="gemini-2.0-flash-lite",  # Adjust model name as needed (e.g., 'xai-model-name')
            messages=[
                {"role": "system", "content": "You are a helpful search engine."},
                {"role": "user", "content": prompt}
            ],
            response_format={"type": "json_object"}
        )
        
        content = response.choices[0].message.content
        results = json.loads(content)
        
        # Handle different possible JSON structures
        if isinstance(results, dict) and "results" in results:
            results = results["results"]
        elif isinstance(results, list):
            pass
        else:
            return None, "Error: Unexpected JSON structure."
        
        return results, None
    
    except Exception as e:
        error_msg = str(e)
        if "404" in error_msg:
            return None, f"Error 404: Model or endpoint not found. Check OPENAI_BASE_URL ({os.environ['OPENAI_BASE_URL']}) and model name."
        elif "401" in error_msg:
            return None, "Error 401: Invalid API key. Check OPENAI_API_KEY."
        else:
            return None, f"Error: {error_msg}"

def display_search_results(query, page=1):
    """Display search results for the given query and page number."""
    results, error = fetch_search_results(query)
    
    if error:
        return error, None, None
    
    # Calculate pagination boundaries
    start_idx = (page - 1) * RESULTS_PER_PAGE
    end_idx = start_idx + RESULTS_PER_PAGE
    total_pages = (len(results) + RESULTS_PER_PAGE - 1) // RESULTS_PER_PAGE
    
    # Ensure indices are within bounds
    if start_idx >= len(results):
        return "No more results to display.", None, None
    
    paginated_results = results[start_idx:end_idx]
    
    # Format results into HTML
    html = """
    <style>
    .search-result {
        margin-bottom: 20px;
    }
    .search-result h3 {
        color: blue;
        font-size: 18px;
        margin: 0;
    }
    .search-result p {
        font-size: 14px;
        margin: 5px 0 0 0;
    }
    .pagination {
        margin-top: 20px;
    }
    </style>
    <div>
    """
    html += f"<h2>Search Results for '{query}' (Page {page} of {total_pages})</h2>"
    html += "<ul>"
    for result in paginated_results:
        title = result.get("title", "No title")
        snippet = result.get("snippet", "No snippet")
        html += f'<li class="search-result"><h3>{title}</h3><p>{snippet}</p></li>'
    html += "</ul>"
    
    # Add pagination controls
    html += '<div class="pagination">'
    if page > 1:
        html += f'<button onclick="update_page({page - 1})">Previous</button>'
    if page < total_pages:
        html += f'<button onclick="update_page({page + 1})">Next</button>'
    html += '</div></div>'
    
    return html, page - 1 if page > 1 else None, page + 1 if page < total_pages else None

def search_handler(query, page):
    """Handle search submission and pagination."""
    html, prev_page, next_page = display_search_results(query, page)
    return html

# Build Gradio interface with Blocks for state management
with gr.Blocks(title="LLM Search Engine") as app:
    gr.Markdown("# LLM Search Engine")
    gr.Markdown("Enter a query below to search using a large language model (press Enter or click Search).")
    
    query_input = gr.Textbox(label="Search Query", placeholder="Type your search here...")
    search_button = gr.Button("Search")
    output_html = gr.HTML()
    
    # Hidden state to track current page
    page_state = gr.State(value=1)
    
    # Define submit behavior
    def on_submit(query, page):
        return search_handler(query, page), page
    
    # Trigger search on Enter key or button click
    query_input.submit(
        fn=on_submit,
        inputs=[query_input, page_state],
        outputs=[output_html, page_state]
    )
    search_button.click(
        fn=on_submit,
        inputs=[query_input, page_state],
        outputs=[output_html, page_state]
    )
    
    # Pagination buttons
    with gr.Row():
        prev_button = gr.Button("Previous", visible=False)
        next_button = gr.Button("Next", visible=False)
    
    def update_page(query, page, direction):
        new_page = page + direction
        html, prev_page, next_page = display_search_results(query, new_page)
        return html, new_page, gr.update(visible=prev_page is not None), gr.update(visible=next_page is not None)
    
    prev_button.click(
        fn=lambda q, p: update_page(q, p, -1),
        inputs=[query_input, page_state],
        outputs=[output_html, page_state, prev_button, next_button]
    )
    
    next_button.click(
        fn=lambda q, p: update_page(q, p, 1),
        inputs=[query_input, page_state],
        outputs=[output_html, page_state, prev_button, next_button]
    )
    
    # Update button visibility after search
    def update_visibility(query, page):
        html, prev_page, next_page = display_search_results(query, page)
        return html, page, gr.update(visible=prev_page is not None), gr.update(visible=next_page is not None)
    
    query_input.submit(
        fn=update_visibility,
        inputs=[query_input, page_state],
        outputs=[output_html, page_state, prev_button, next_button]
    )
    search_button.click(
        fn=update_visibility,
        inputs=[query_input, page_state],
        outputs=[output_html, page_state, prev_button, next_button]
    )

app.launch()