import os
import logging
import gradio as gr
from huggingface_hub import InferenceClient
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.schema import Document
from duckduckgo_search import DDGS

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

MODELS = [
    "mistralai/Mistral-7B-Instruct-v0.3",
    "mistralai/Mixtral-8x7B-Instruct-v0.1",
    "mistralai/Mistral-Nemo-Instruct-2407"
]

def get_embeddings():
    return HuggingFaceEmbeddings(model_name="sentence-transformers/stsb-roberta-large")

def duckduckgo_search(query):
    with DDGS() as ddgs:
        results = ddgs.text(query, max_results=5)
    return results

def create_web_search_vectors(search_results):
    embed = get_embeddings()
    documents = []
    for result in search_results:
        if 'body' in result:
            content = f"{result['title']}\n{result['body']}\nSource: {result['href']}"
            documents.append(Document(page_content=content, metadata={"source": result['href']}))
    return FAISS.from_documents(documents, embed)

def get_response_with_search(query, model, num_calls=3, temperature=0.2):
    search_results = duckduckgo_search(query)
    web_search_database = create_web_search_vectors(search_results)
    if not web_search_database:
        yield "No web search results available. Please try again.", ""
        return

    retriever = web_search_database.as_retriever(search_kwargs={"k": 5})
    relevant_docs = retriever.get_relevant_documents(query)
    context = "\n".join([doc.page_content for doc in relevant_docs])
    prompt = f"""Using the following context from web search results:
{context}
Write a detailed and complete research document that fulfills the following user request: '{query}'
After writing the document, please provide a list of sources used in your response."""

    # Use Hugging Face API
    client = InferenceClient(model, token=huggingface_token)
    main_content = ""
    for i in range(num_calls):
        for message in client.chat_completion(
            messages=[{"role": "user", "content": prompt}],
            max_tokens=10000,
            temperature=temperature,
            stream=True,
        ):
            if message.choices and message.choices[0].delta and message.choices[0].delta.content:
                chunk = message.choices[0].delta.content
                main_content += chunk
                yield main_content, ""

def respond(message, history, model, temperature, num_calls):
    logging.info(f"User Query: {message}")
    logging.info(f"Model Used: {model}")
    logging.info(f"Temperature: {temperature}")
    logging.info(f"Number of API Calls: {num_calls}")

    try:
        for main_content, sources in get_response_with_search(message, model, num_calls=num_calls, temperature=temperature):
            response = f"{main_content}\n\n{sources}"
            yield response
    except Exception as e:
        logging.error(f"Error in respond function: {str(e)}")
        yield f"An error occurred: {str(e)}"

css = """
/* Fine-tune chatbox size */
.chatbot-container {
    height: 600px !important;
    width: 100% !important;
}
.chatbot-container > div {
    height: 100%;
    width: 100%;
}
"""

# Gradio interface setup
def create_gradio_interface():
    custom_placeholder = "Enter your question here for web search."

    demo = gr.ChatInterface(
        respond,
        additional_inputs=[
            gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[2]),
            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 Assistant",
        description="Use web search to answer questions or generate summaries.",
        theme=gr.Theme.from_hub("allenai/gradio-theme"),
        css=css,
        examples=[
            ["What are the latest developments in artificial intelligence?"],
            ["Explain the concept of quantum computing."],
            ["What are the environmental impacts of renewable energy?"]
        ],
        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,
        )
    )

    with demo:
        gr.Markdown("""
        ## How to use
        1. Enter your question in the chat interface.
        2. Select the model you want to use from the dropdown.
        3. Adjust the Temperature to control the randomness of the response.
        4. Set the Number of API Calls to determine how many times the model will be queried.
        5. Press Enter or click the submit button to get your answer.
        6. Use the provided examples or ask your own questions.
        """)

    return demo

if __name__ == "__main__":
    demo = create_gradio_interface()
    demo.launch(share=True)