import gradio as gr
from huggingface_hub import InferenceClient
import requests
from bs4 import BeautifulSoup
from bs4.element import Comment

def tag_visible(element):
    if element.parent.name in ['style', 'script', 'head', 'title', 'meta', '[document]']:
        return False
    if isinstance(element, Comment):
        return False
    return True

def get_text_from_url(url):
    response = requests.get(url)
    soup = BeautifulSoup(response.text, 'html.parser')
    texts = soup.find_all(text=True)
    visible_texts = filter(tag_visible, texts)
    return "\n".join(t.strip() for t in visible_texts)

# Gather text from your homepage (and any extensions)
text_list = []
homepage_url = "https://sites.google.com/view/abhilashnandy/home/"
extensions = ["", "pmrf-profile-page"]
for ext in extensions:
    url_text = get_text_from_url(homepage_url + ext)
    text_list.append(url_text)

# Build the system message with homepage info.
SYSTEM_MESSAGE = (
    "You are a QA chatbot to answer queries (in less than 30 words) on my homepage that has the following information -\n\n"
    + "\n\n".join(text_list)
    + "\n\n"
)

# Use the GPTQ variant of TinyLlama which includes the tokenizer configuration
client = InferenceClient("TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ")

def respond(message, history: list[tuple[str, str]], system_message=SYSTEM_MESSAGE,
            max_tokens=140, temperature=0.7, top_p=0.95):
    messages = [{"role": "system", "content": system_message}]
    for val in history:
        if len(val) >= 1:
            messages.append({"role": "user", "content": "Question: " + val[0]})
        if len(val) >= 2:
            messages.append({"role": "assistant", "content": "Answer: " + val[1]})
    messages.append({"role": "user", "content": message})
    try:
        response = client.chat_completion(
            messages,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            # stream=True,  # Uncomment this if you want streaming output for debugging.
        )
        return response.choices[0].message["content"]
    except Exception as e:
        print(f"An error occurred: {e}")
        return str(e)

initial_message = [("user", "Yo who dis Abhilash?")]
markdown_note = "## Ask Anything About Me! (Might show a tad bit of hallucination!)"

demo = gr.Blocks()
with demo:
    gr.Markdown(markdown_note)
    gr.ChatInterface(
        fn=respond,
        # examples=["Yo who dis Abhilash?", "What is Abhilash's most recent publication?"],
        additional_inputs=[
            # Additional components can be added here if needed.
        ],
    )

if __name__ == "__main__":
    demo.launch()