import gradio as gr
from huggingface_hub import InferenceClient
from bs4 import BeautifulSoup
import requests
import os
from typing import List, Tuple, Optional

# Retrieve API key from environment variable
api_key = os.getenv('HFTOKEN')
if not api_key:
    raise ValueError("API key not found. Please set the HFTOKEN environment variable.")

# Initialize the InferenceClient with the specified model and API key
client = InferenceClient(
    model="meta-llama/Meta-Llama-3.1-405B-Instruct",
    token=api_key
)

def scrape_yahoo_search(query: str, num_results: int = 3) -> Tuple[str, str]:
    """
    Scrapes Yahoo search results for the given query and returns detailed snippets and URLs for the top results.
    
    Args:
        query (str): The search query.
        num_results (int): Number of results to retrieve.
    
    Returns:
        Tuple[str, str]: The formatted snippets and URLs of the top search results.
    """
    search_url = f"https://search.yahoo.com/search?p={query}"
    headers = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
    }
    
    try:
        response = requests.get(search_url, headers=headers)
        response.raise_for_status()
        soup = BeautifulSoup(response.content, 'html.parser')
        
        results = []
        result_elements = soup.find_all('div', {'class': 'dd algo algo-sr Sr'}, limit=num_results)
        
        if result_elements:
            for result in result_elements:
                title = result.find('h3').get_text(strip=True) if result.find('h3') else "No title"
                snippet = result.find('div', {'class': 'compText aAbs'}).get_text(strip=True) if result.find('div', {'class': 'compText aAbs'}) else "No snippet"
                url = result.find('a')['href'] if result.find('a') else "No URL"
                results.append((title, snippet, url))
            
            formatted_results = "\n\n".join(
                f"Title: {title}\nSnippet: {snippet}\nURL: {url}" for title, snippet, url in results
            )
            return formatted_results, search_url
        else:
            return "No results found.", search_url
    
    except requests.RequestException as e:
        return f"Request error: {str(e)}", search_url
    except Exception as e:
        return f"Processing error: {str(e)}", search_url

def extract_search_query(message: str, trigger_word: str) -> Optional[str]:
    """
    Extracts the search query from the message based on the trigger word.
    
    Args:
        message (str): The user's input message.
        trigger_word (str): The word that activates the search feature.
    
    Returns:
        Optional[str]: The extracted search query if found, otherwise None.
    """
    lower_message = message.lower()
    if trigger_word in lower_message:
        parts = lower_message.split(trigger_word, 1)
        if len(parts) > 1:
            query = parts[1].strip()
            return query if query else None
    return None

def respond(
    message: str,
    history: List[Tuple[str, str]],
    system_message: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
) -> str:
    """
    Generates a response from the AI model based on the user's message, chat history, and optional Yahoo search results.
    
    Args:
        message (str): The user's input message.
        history (List[Tuple[str, str]]): A list of tuples representing the conversation history (user, assistant).
        system_message (str): A system-level message guiding the AI's behavior.
        max_tokens (int): The maximum number of tokens for the output.
        temperature (float): Sampling temperature for controlling the randomness.
        top_p (float): Top-p (nucleus sampling) for controlling diversity.
    
    Returns:
        str: The AI's response as it is generated, including the source URL if applicable.
    """
    
    # Check for trigger word and activate search feature if present
    trigger_word = "search"
    query = extract_search_query(message, trigger_word)
    
    if query:
        snippet, url = scrape_yahoo_search(query, num_results=3)
        message += f"\n\nWeb Content:\n{snippet}\nSource: {url}"
    elif query is None:
        message = "Please provide a search query after the trigger word."

    # Prepare the conversation history for the API call
    messages = [{"role": "system", "content": system_message}]
    
    for user_input, assistant_response in history:
        if user_input:
            messages.append({"role": "user", "content": user_input})
        if assistant_response:
            messages.append({"role": "assistant", "content": assistant_response})

    # Add the latest user message to the conversation
    messages.append({"role": "user", "content": message})

    # Initialize an empty response
    response = ""

    try:
        # Generate a response from the model with streaming
        for response_chunk in client.chat_completion(
            messages=messages,
            max_tokens=max_tokens,
            stream=True,
            temperature=temperature,
            top_p=top_p,
        ):
            token = response_chunk.choices[0].delta.content
            response += token
    except Exception as e:
        return f"AI model error: {str(e)}"

    return response

# Define the ChatInterface with additional input components for user customization
demo = gr.ChatInterface(
    fn=respond,
    additional_inputs=[
        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ],
    title="Chatbot Interface",
    description="A customizable chatbot interface using Hugging Face's Inference API with Yahoo search scraping capabilities.",
)

# Launch the Gradio interface
if __name__ == "__main__":
    demo.launch()