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

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
        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>'

        # Generate choices for dropdown
        choices = [f"{row['Topic']} ({row['Date']})" for row in data]
        return html, choices
        
    except Exception as e:
        print(f"Error in extract_table: {e}")
        return f"<p>Error: {str(e)}</p>", []

def prepare_topic(selected_topic):
    print(f"Preparing topic: {selected_topic}")  # Debug print
    try:
        if not selected_topic:
            return "Please select a topic first"
            
        # Handle potential list or string input
        if isinstance(selected_topic, list):
            selected_topic = selected_topic[0] if selected_topic else ""
        
        # Find the index of the selected topic
        for row in data:
            full_topic = f"{row['Topic']} ({row['Date']})"
            if full_topic == selected_topic:
                topic = row["Topic"]
                date = row["Date"]
                message = f"Please prepare a 15-minutes reading material covering main topics for '{topic}' lecture scheduled for {date}"
                print(f"Generated preparation message: {message}")  # Debug print
                return message
        
        print(f"Topic not found: {selected_topic}")
        return "Error: Topic not found"
    
    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):
            # Dropdown for selecting topic
            topic_dropdown = gr.Dropdown(
                label="Select Topic",
                choices=[],
                interactive=True,
                value=None
            )

            prepare_btn = gr.Button("Prepare Topic")
            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")
            
            
            
        with gr.Column(scale=3):
            chatbot = gr.Chatbot()
            msg = gr.Textbox(label="Message")
            system_message = gr.Textbox(
                value="Students lecture preparation companion.",
                label="System message"
            )

            with gr.Row():
                submit = gr.Button("Submit")
                clear = gr.Button("Clear")
    
    # Event handlers
    # Extract table and update dropdown
    def update_interface(url):
        html, choices = extract_table(url)
        return html, gr.Dropdown(choices=choices)

    extract_btn.click(
        fn=update_interface,
        inputs=[url_input],
        outputs=[table_output, 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]
    )
    
    # Message submit handlers
    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(share=True)