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):
    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=2048,
        stream=True,
        temperature=0.7,
        top_p=0.9,
    ):
        if message.choices[0].delta.content is not None:
            response += message.choices[0].delta.content
            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 = '''
        <style>
            .dataframe {
                border-collapse: collapse;
                width: 100%;
                margin: 10px 0;
            }
            .dataframe th, .dataframe td {
                border: 1px solid #ddd;
                padding: 8px;
                text-align: left;
            }
            .dataframe th {
                background-color: #f6f8fa;
            }
            .dataframe tr:nth-child(even) {
                background-color: #f9f9f9;
            }
        </style>
        '''
        
        html += '<table class="dataframe">'
        html += '<thead><tr><th>Date</th><th>Topic</th></tr></thead>'
        html += '<tbody>'
        
        for row in data:
            html += f'''
                <tr>
                    <td>{row['Date']}</td>
                    <td>{row['Topic']}</td>
                </tr>
            '''
        html += '</tbody></table>'
        return html
        
    except Exception as e:
        print(f"Error in extract_table: {e}")  # Debug print
        return f"<p>Error: {str(e)}</p>"

def prepare_topic(index):
    print(f"Preparing topic with index: {index}")  # Debug print
    try:
        # Convert index to integer
        index = int(index)
        
        # Check if index is valid
        if 0 <= index < len(data):
            topic = data[index]["Topic"]
            date = data[index]["Date"]
            message = f"Please prepare a 10-minute reading guide for the topic '{topic}' scheduled for {date}"
            print(f"Generated preparation message: {message}")  # Debug print
            return message
        else:
            print(f"Invalid index: {index}. Total data length: {len(data)}")
            return "Error: Invalid topic selection"
    except ValueError:
        print(f"Could not convert {index} to integer")
        return "Error: Invalid topic selection"
    except Exception as e:
        print(f"Unexpected error in prepare_topic: {e}")
        return "Error: Could not prepare topic"

def add_text(history, text):
    history = history + [(text, None)]
    return history

def generate_response(history, system_message):
    if not history:
        return history
    
    response = ""
    for chunk in respond(history[-1][0], history[:-1], system_message):
        response = chunk
        history[-1] = (history[-1][0], response)
        yield history

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")
            
            # Dropdown for selecting topic
            topic_dropdown = gr.Dropdown(
                label="Select Topic",
                choices=[],
                interactive=True
            )
            prepare_btn = gr.Button("Prepare Topic")
            
        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]
    ).success(
        fn=lambda: [f"{row['Topic']} ({row['Date']})" for row in data],
        outputs=[topic_dropdown]
    )
    
    # Prepare topic handler
    prepare_btn.click(
        fn=prepare_topic,
        inputs=[topic_dropdown],
        outputs=[msg]
    ).success(
        fn=add_text,
        inputs=[chatbot, msg],
        outputs=[chatbot],
        queue=False
    ).then(
        fn=generate_response,
        inputs=[chatbot, system_message],
        outputs=[chatbot]
    )
    
    # Rest of your event handlers remain the same
    msg.submit(
        fn=add_text,
        inputs=[chatbot, msg],
        outputs=[chatbot]
    ).success(
        fn=lambda: "",
        outputs=[msg]
    ).then(
        fn=generate_response,
        inputs=[chatbot, system_message],
        outputs=[chatbot]
    )
    
    submit.click(
        fn=add_text,
        inputs=[chatbot, msg],
        outputs=[chatbot]
    ).success(
        fn=lambda: "",
        outputs=[msg]
    ).then(
        fn=generate_response,
        inputs=[chatbot, system_message],
        outputs=[chatbot]
    )
    
    # Clear button handler
    clear.click(fn=clear_chat, outputs=[chatbot, msg])

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