import gradio as gr
from duckduckgo_search import DDGS
from typing import List, Dict
import os
import logging

logging.basicConfig(level=logging.INFO)

# Environment variables and configurations
huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")

class ConversationManager:
    def __init__(self):
        self.history = []
        self.current_context = None

    def add_interaction(self, query, response):
        self.history.append((query, response))
        self.current_context = f"Previous query: {query}\nPrevious response summary: {response[:200]}..."

    def get_context(self):
        return self.current_context

def get_web_search_results(query: str, max_results: int = 10) -> List[Dict[str, str]]:
    try:
        results = list(DDGS().text(query, max_results=max_results))
        if not results:
            print(f"No results found for query: {query}")
        return results
    except Exception as e:
        print(f"An error occurred during web search: {str(e)}")
        return [{"error": f"An error occurred during web search: {str(e)}"}]

def rephrase_query(original_query: str, conversation_manager: ConversationManager) -> str:
    context = conversation_manager.get_context()
    if context:
        prompt = f"""You are a highly intelligent conversational chatbot. Your task is to analyze the given context and new query, then decide whether to rephrase the query with or without incorporating the context. Follow these steps:

        1. Determine if the new query is a continuation of the previous conversation or an entirely new topic.
        2. If it's a continuation, rephrase the query by incorporating relevant information from the context to make it more specific and contextual.
        3. If it's a new topic, rephrase the query to make it more appropriate for a web search, focusing on clarity and accuracy without using the previous context.
        4. Provide ONLY the rephrased query without any additional explanation or reasoning.
        
        Context: {context}
        
        New query: {original_query}
        
        Rephrased query:"""
        response = DDGS().chat(prompt, model="llama-3.1-70b")
        # Extract only the rephrased query, removing any explanations
        rephrased_query = response.split('\n')[0].strip()
        return rephrased_query
    return original_query

def summarize_results(query: str, search_results: List[Dict[str, str]], conversation_manager: ConversationManager) -> str:
    try:
        context = conversation_manager.get_context()
        search_context = "\n\n".join([f"Title: {result['title']}\nContent: {result['body']}" for result in search_results])

        prompt = f"""You are a highly intelligent & expert analyst and your job is to skillfully articulate the web search results about '{query}' and considering the context: {context}, 
        You have to create a comprehensive news summary FOCUSING on the context provided to you. 
        Include key facts, relevant statistics, and expert opinions if available. 
        Ensure the article is well-structured with an introduction, main body, and conclusion, IF NECESSARY. 
        Address the query in the context of the ongoing conversation IF APPLICABLE.
        Cite sources directly within the generated text and not at the end of the generated text, integrating URLs where appropriate to support the information provided:

        {search_context}

        Article:"""

        summary = DDGS().chat(prompt, model="llama-3-70b")
        return summary
    except Exception as e:
        return f"An error occurred during summarization: {str(e)}"

conversation_manager = ConversationManager()

def respond(message, chat_history, temperature, num_api_calls):
    final_summary = ""
    original_query = message
    rephrased_query = rephrase_query(message, conversation_manager)

    logging.info(f"Original query: {original_query}")
    logging.info(f"Rephrased query: {rephrased_query}")

    for _ in range(num_api_calls):
        search_results = get_web_search_results(rephrased_query)

        if not search_results:
            final_summary += f"No search results found for the query: {rephrased_query}\n\n"
        elif "error" in search_results[0]:
            final_summary += search_results[0]["error"] + "\n\n"
        else:
            summary = summarize_results(rephrased_query, search_results, conversation_manager)
            final_summary += summary + "\n\n"

    if final_summary:
        conversation_manager.add_interaction(original_query, final_summary)
        return final_summary
    else:
        return "Unable to generate a response. Please try a different query."

# The rest of your code (CSS, theme, and Gradio interface setup) remains the same
css = """
Your custom CSS here
"""

custom_placeholder = "Ask me anything about web content"

theme = gr.themes.Soft(
    primary_hue="orange",
    secondary_hue="amber",
    neutral_hue="gray",
    font=[gr.themes.GoogleFont("Exo"), "ui-sans-serif", "system-ui", "sans-serif"]
).set(
    body_background_fill_dark="#0c0505",
    block_background_fill_dark="#0c0505",
    block_border_width="1px",
    block_title_background_fill_dark="#1b0f0f",
    input_background_fill_dark="#140b0b",
    button_secondary_background_fill_dark="#140b0b",
    border_color_accent_dark="#1b0f0f",
    border_color_primary_dark="#1b0f0f",
    background_fill_secondary_dark="#0c0505",
    color_accent_soft_dark="transparent",
    code_background_fill_dark="#140b0b"
)

demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Slider(minimum=0.1, maximum=1.0, value=0.2, step=0.1, label="Temperature"),
        gr.Slider(minimum=1, maximum=5, value=1, step=1, label="Number of API Calls")
    ],
    title="AI-powered Web Search and PDF Chat Assistant",
    description="This AI-powered Web Search and PDF Chat Assistant combines real-time web search capabilities with advanced language processing.",
    theme=theme,
    css=css,
    examples=[
        ["What is AI"],
        ["Any recent news on US Banks"],
        ["Who is Donald Trump"]
    ],
    cache_examples=False,
    analytics_enabled=False,
    textbox=gr.Textbox(placeholder=custom_placeholder, container=False, scale=7),
    chatbot=gr.Chatbot(
        show_copy_button=True,
        likeable=True,
        layout="bubble",
        height=400,
    )
)

demo.launch()