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

# Initialize HF client
client = InferenceClient("meta-llama/Llama-2-7b-chat-hf")

def respond(message, history, max_tokens=512, temperature=0.7, top_p=0.95):
    try:
        # Format messages including history
        messages = []
        for user_msg, assistant_msg in history:
            messages.append({"role": "user", "content": user_msg})
            messages.append({"role": "assistant", "content": assistant_msg})
        messages.append({"role": "user", "content": message})
        
        # Generate response
        response = ""
        for chunk in client.chat_completion(
            messages,
            max_tokens=max_tokens,  # Changed back to max_tokens
            temperature=temperature,
            top_p=top_p,
            stream=True,
        ):
            if hasattr(chunk.choices[0].delta, 'content'):
                token = chunk.choices[0].delta.content
                if token:
                    response += token
        return response
        
    except Exception as e:
        return f"Error: {str(e)}"
    
def extract_schedule(url):
    try:
        # Fetch and parse webpage
        response = requests.get(url)
        response.raise_for_status()
        soup = BeautifulSoup(response.text, 'html.parser')
        
        # Find table and extract data
        table = soup.find('table')
        if not table:
            return "<p>No table found on page</p>"
            
        schedule_data = []
        rows = table.find_all('tr')
        
        for row in rows[1:]:  # Skip header row
            cells = row.find_all('td')
            if len(cells) >= 4:  # Only process rows with enough columns
                date = cells[0].text.strip()
                topic = cells[1].text.strip()
                
                # Skip empty rows and non-lecture entries
                if date and topic and not topic.startswith('See Canvas'):
                    schedule_data.append({
                        'Date': date[:10],  # Only YYYY-MM-DD
                        'Topic': topic
                    })
        
        # Create DataFrame
        df = pd.DataFrame(schedule_data)
        
        # Convert to HTML with styling
        html = f"""
        <style>
        table {{ border-collapse: collapse; width: 100%; }}
        th, td {{ 
            border: 1px solid black;
            padding: 8px;
            text-align: left;
        }}
        th {{ background-color: #f2f2f2; }}
        </style>
        {df.to_html(index=False)}
        """
        return html
        
    except Exception as e:
        return f"<p>Error: {str(e)}</p>"

def display_schedule(url):
    try:
        html_table = extract_schedule(url)
        return html_table  # Already HTML string
    except Exception as e:
        return str(e)

with gr.Blocks() as demo:
    with gr.Row():
        # Left Column - Schedule
        with gr.Column(scale=1):
            url_input = gr.Textbox(
                value="https://id2223kth.github.io/schedule/",
                label="Schedule URL"
            )
            schedule_output = gr.HTML(label="Extracted Schedule")
            extract_btn = gr.Button("Extract Schedule")
            
            extract_btn.click(
                fn=display_schedule,
                inputs=[url_input],
                outputs=[schedule_output]
            )
        
        # Right Column - Chatbot
        with gr.Column(scale=2):
            chatbot = gr.ChatInterface(
                respond,
                #additional_inputs=[
                #    gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
                #    gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
                #    gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
                #]
            )

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