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 using inline JavaScript
        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="(() => {{
                                const index = document.getElementById('prepare-index');
                                index.value = '{i}';
                                index.dispatchEvent(new Event('input'));
                                document.getElementById('prepare-trigger').click();
                            }})();"
                            style="padding: 5px 10px; cursor: pointer; background-color: #0366d6; color: white; border: none; border-radius: 4px;"
                        >
                            Prepare
                        </button>
                    </td>
                </tr>
            '''
        html += '</tbody></table>'
        
        # Add some CSS to make the table look better
        html = f'''
            <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;
                }}
                .dataframe button:hover {{
                    background-color: #0056b3 !important;
                }}
            </style>
            {html}
        '''
        return html
    except Exception as e:
        return f"<p>Error: {str(e)}</p>"

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 prepare_topic_message(index):
    try:
        print(f"Received index: {index}")  # Debug print
        
        # Handle the index as a number
        idx = int(index)
        if 0 <= idx < len(data):
            topic = data[idx]["Topic"]
            date = data[idx]["Date"]
            message = f"Please prepare a 10-minute reading guide for the topic '{topic}' scheduled for {date}"
            print(f"Generated message: {message}")  # Debug print
            return message
        else:
            print(f"Index {idx} out of range")
            return ""
    except Exception as e:
        print(f"Error in prepare_topic_message: {e}")
        return ""

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_index = gr.Textbox(value="", visible=False, elem_id="prepare-index")
            prepare_trigger = gr.Button("Prepare", visible=False, elem_id="prepare-trigger")
            
        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]
    )
    
    # Prepare trigger handler
    # Prepare trigger handler
    prepare_trigger.click(
        fn=prepare_topic_message,
        inputs=prepare_index,
        outputs=msg,
        queue=False  # Changed from api_name to queue=False
    ).success(
        fn=add_text,
        inputs=[chatbot, msg],
        outputs=chatbot,
        queue=False
    ).then(
        fn=generate_response,
        inputs=[chatbot, system_message],
        outputs=chatbot
    )
    
    # Submit button handler
    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()