import os
import json
import gradio as gr
from openai import OpenAI
import requests
import time
import google.generativeai as genai
from groq import Groq
from urllib.parse import urljoin
from bs4 import BeautifulSoup
import re
import logging
from groq import Groq

# Load API keys from environment variables
HF_API_TOKEN = os.getenv("HF_API_KEY")
DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")

# Initialize response time records
response_times = {"Gemini Flash 2.0":[], "Deepseek":[]}

# Initialize Deepseek client
Deepseek_client = OpenAI(api_key=DEEPSEEK_API_KEY, base_url="https://api.deepseek.com")

def scrape_website(url):
    """
    Scrapes a website and its subpages, respecting robots.txt.
    """
    url_text = []
    try:
        response = requests.get(url, timeout=10)
        response.raise_for_status()

        soup = BeautifulSoup(response.content, "html.parser")
        links = []
        for link in soup.find_all('a', href=True):
            absolute_url = urljoin(url, link['href'])
            links.append(absolute_url)

        # Basic filtering (improve this as needed)
        valid_links = []
        for link in links:
            if re.match(r'^https?://', link) and link.startswith(url): #same domain links only
              valid_links.append(link)
        
        print(f"Found {len(valid_links)} links on {url}")
        
        for link in valid_links:
            print(f"Scraping {link}")
            # Add logic to check robots.txt for the subpage (optional)
            try:
                subpage_response = requests.get(link, timeout=10)
                subpage_response.raise_for_status()
                # Process the subpage content here. For example:
                subpage_soup = BeautifulSoup(subpage_response.content, "html.parser")
                url_text.append(subpage_soup.get_text())
                print(f"Title: {subpage_soup.title.string if subpage_soup.title else 'No Title'}")
                # Extract relevant information from subpage_soup...
            except requests.exceptions.RequestException as e:
                print(f"Error scraping {link}: {e}")
        
    except requests.exceptions.RequestException as e:
        print(f"Error accessing {url}: {e}")
    return url_text

def check_prime(n):
    if n < 2:
        return False
    for i in range(2, int(n**0.5) + 1):
        if n % i == 0:
            return False
    return True

def query_deepseek(prompt, chatbot_state):
    """Queries Deepseek V3 model with the full conversation history as context."""
    start_time = time.time()

    response = Deepseek_client.chat.completions.create(
        model="deepseek-chat",
        messages=[
            {"role": "system", "content": "You are a helpful assistant"},
            {"role": "user", "content": prompt},
        ],
        stream=False
    )
    response_times["Deepseek"].append(time.time() - start_time)
    return response.choices[0].message.content

def query_gemini(prompt, chatbot_state):
    """Queries Google Gemini."""
    instruction = "You are a helpful system bot."
    start_time = time.time()
    try:
        model = genai.GenerativeModel(
        model_name='gemini-2.0-flash-exp',
        tools='code_execution')

        # Format conversation history
        history = " ".join([f"{entry['role'].capitalize()}: {entry['content']}" for entry in chatbot_state])
        full_prompt = f"{history}\nUser: {prompt}"
        response = model.generate_content(full_prompt)
        return response.text
    except Exception as e:
        return f"Gemini Error: {str(e)}"
    finally:
        response_times["Gemini Flash 2.0"].append(time.time() - start_time)


def chatbot_response(prompt, model_choice, chatbot_state):
    """Handles the conversation flow and updates the chatbot state."""
    # Add user message to the conversation history
    chatbot_state.append({"role": "user", "content": prompt})

    # Generate a response based on the selected model
    if model_choice == "Gemini Flash 2.0":
        response = query_gemini(prompt, chatbot_state)
    elif model_choice == "Deepseek":
        response = query_deepseek(prompt, chatbot_state)
    else:
        response = "Invalid model choice."

    # Add the AI's response to the conversation history
    chatbot_state.append({"role": "assistant", "content": response})

    return chatbot_state, ""

def reset_memory():
    global chatbot_state
    chatbot_state = gr.State([])
    return [], "", []

def get_average_times():
    """Calculate and display average response times for all models."""
    avg_times = {model: (sum(times) / len(times) if times else 0) for model, times in response_times.items()}
    return f"Average Response Times (in seconds):\n" + "\n".join([f"{model}: {time:.2f}s" for model, time in avg_times.items()])

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("## Inference Latency Comparison")
    
    # Model selection
    model_choice = gr.Radio(["Gemini Flash 2.0", "Deepseek"], label="Choose AI Model", value="Deepseek")
    
    # Chatbot component
    chatbot = gr.Chatbot(label="Conversation", type="messages")
    chatbot_state = gr.State([])  # Shared conversation history
    
    # User input
    chatbot_input = gr.Textbox(label="Enter your message:", placeholder="Type your message here")
    
    # Submit button
    submit_button = gr.Button("Submit")
    submit_button.click(
        fn=chatbot_response,
        inputs=[chatbot_input, model_choice, chatbot_state],
        outputs=[chatbot, chatbot_input]
    )

    # Get times button
    output_box = gr.Textbox(label="Average Response Times", lines=10)
    get_times_button = gr.Button("Average Times")
    
    # Connect button to the function
    get_times_button.click(
        fn=get_average_times,
        inputs=[],  # No inputs required
        outputs=output_box  # Display results in the output textbox
    )
    
    # Reset button
    reset_button = gr.Button("Reset Memory")

    reset_button.click(
    fn=reset_memory,  # Ensure the reset function is called
    inputs=[],        # No inputs needed for reset
    outputs=[chatbot, chatbot_input, chatbot_state]  # Clear the chatbot UI and the state
)

demo.launch(auth=("latham", "1956"))