import gradio as gr
from huggingface_hub import InferenceClient
import requests
from bs4 import BeautifulSoup
import pandas as pd

client = InferenceClient("meta-llama/Llama-3.2-3B-Instruct")

# Global data store for the table
data = []

def respond(message, history, system_message, max_tokens=2048, temperature=0.7, top_p=0.9):
    messages = [{"role": "system", "content": system_message}]
    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})
    messages.append({"role": "user", "content": message})
    response = ""
    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = message.choices[0].delta.content
        response += token
        yield response

def extract_table(url):
    global data
    try:
        response = requests.get(url)
        response.raise_for_status()
        soup = BeautifulSoup(response.text, 'html.parser')
        table = soup.find('table')
        if not table:
            return "<p>No table found on page</p>"
        
        # Clear existing data
        data = []
        rows = table.find_all('tr')
        for i, row in enumerate(rows[1:]):
            cells = row.find_all('td')
            if len(cells) >= 2:
                data.append({
                    'Date': cells[0].text.strip()[:10],
                    'Topic': cells[1].text.strip(),
                })
        
        # Create HTML table with prepare buttons
        html = '<table class="dataframe">'
        html += '<thead><tr><th>Date</th><th>Topic</th><th>Action</th></tr></thead>'
        html += '<tbody>'
        for i, row in enumerate(data):
            html += f'''
                <tr>
                    <td>{row['Date']}</td>
                    <td>{row['Topic']}</td>
                    <td>
                        <button onclick='
                            document.getElementById("prepare-topic").value = "{i}";
                            document.getElementById("prepare-topic").dispatchEvent(new Event("input"));
                            document.getElementById("prepare-button").click();
                        '>
                            Prepare
                        </button>
                    </td>
                </tr>
            '''
        html += '</tbody></table>'
        return html
    except Exception as e:
        return f"<p>Error: {str(e)}</p>"

def user_message(message, history):
    history = history or []
    history.append((message, None))
    return "", history

def handle_prepare(index, history):
    try:
        index = int(index)
        if 0 <= index < len(data):
            topic = data[index]["Topic"]
            message = f"Prepare a 10-minute reading for the topic: {topic}"
            history = history or []
            history.append((message, None))
            return history
    except:
        return history or []

def clear_chat():
    return [], ""

# Gradio app
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column(scale=1):
            url_input = gr.Textbox(
                value="https://id2223kth.github.io/schedule/",
                label="Table URL"
            )
            table_output = gr.HTML(label="Extracted Table")
            extract_btn = gr.Button("Extract Table")
            
            # Hidden components for prepare functionality
            prepare_topic = gr.Textbox(value="", visible=False, elem_id="prepare-topic")
            prepare_button = gr.Button("Prepare", visible=False, elem_id="prepare-button")
            
        with gr.Column(scale=2):
            chatbot = gr.Chatbot()
            msg = gr.Textbox(label="Message")
            system_message = gr.Textbox(
                value="Student class preparation companion.",
                label="System message"
            )

            with gr.Row():
                submit = gr.Button("Submit")
                clear = gr.Button("Clear")
    
    # Event handlers
    extract_btn.click(
        fn=extract_table,
        inputs=[url_input],
        outputs=[table_output]
    )
    
    # Submit button handler
    submit.click(
        fn=user_message,
        inputs=[msg, chatbot],
        outputs=[msg, chatbot]
    ).then(
        fn=respond,
        inputs=[msg, chatbot, system_message],
        outputs=[chatbot]
    )
    
    # Prepare button handler
    prepare_button.click(
        fn=handle_prepare,
        inputs=[prepare_topic, chatbot],
        outputs=[chatbot]
    ).then(
        fn=respond,
        inputs=[prepare_topic, chatbot, system_message],
        outputs=[chatbot]
    )
    
    # Clear button handler
    clear.click(fn=clear_chat, outputs=[chatbot, msg])

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