import gradio as gr
from PyPDF2 import PdfReader
from bs4 import BeautifulSoup
import openai
import traceback
import requests
from io import BytesIO
from transformers import AutoTokenizer
import json
from datetime import datetime
import os
from openai import OpenAI
import re

# Cache for tokenizers to avoid reloading
tokenizer_cache = {}

# Global variables for providers
PROVIDERS = {
    "SambaNova": {
        "name": "SambaNova",
        "logo": "https://venturebeat.com/wp-content/uploads/2020/02/SambaNovaLogo_H_F.jpg",
        "endpoint": "https://api.sambanova.ai/v1/",
        "api_key_env_var": "SAMBANOVA_API_KEY",
        "models": [
            "Meta-Llama-3.1-70B-Instruct",
        ],
        "type": "tuples",
        "max_total_tokens": "50000",
    },
    "Hyperbolic": {
        "name": "hyperbolic",
        "logo": "https://www.nftgators.com/wp-content/uploads/2024/07/Hyperbolic.jpg",
        "endpoint": "https://api.hyperbolic.xyz/v1",
        "api_key_env_var": "HYPERBOLIC_API_KEY",
        "models": [
            "meta-llama/Llama-3.3-70B-Instruct",
            "meta-llama/Meta-Llama-3.1-405B-Instruct",
        ],
        "type": "tuples",
        "max_total_tokens": "50000",
    },
}


# Functions for paper fetching
def fetch_paper_info_neurips(paper_id):
    url = f"https://openreview.net/forum?id={paper_id}"
    response = requests.get(url)
    if response.status_code != 200:
        return None, None, None

    html_content = response.content
    soup = BeautifulSoup(html_content, 'html.parser')

    # Extract title
    title_tag = soup.find('h2', class_='citation_title')
    title = title_tag.get_text(strip=True) if title_tag else 'Title not found'

    # Extract authors
    authors = []
    author_div = soup.find('div', class_='forum-authors')
    if author_div:
        author_tags = author_div.find_all('a')
        authors = [tag.get_text(strip=True) for tag in author_tags]
    author_list = ', '.join(authors) if authors else 'Authors not found'

    # Extract abstract
    abstract_div = soup.find('strong', text='Abstract:')
    if abstract_div:
        abstract_paragraph = abstract_div.find_next_sibling('div')
        abstract = abstract_paragraph.get_text(strip=True) if abstract_paragraph else 'Abstract not found'
    else:
        abstract = 'Abstract not found'

    # Construct preamble
    link = f"https://openreview.net/forum?id={paper_id}"
    return title, author_list, f"**Abstract:** {abstract}\n\n[View on OpenReview]({link})"


def fetch_paper_content_neurips(paper_id):
    try:
        url = f"https://openreview.net/pdf?id={paper_id}"
        response = requests.get(url)
        response.raise_for_status()
        pdf_content = BytesIO(response.content)
        reader = PdfReader(pdf_content)
        text = ""
        for page in reader.pages:
            text += page.extract_text()
        return text
    except:
        return None


def fetch_paper_content_arxiv(paper_id):
    try:
        url = f"https://arxiv.org/pdf/{paper_id}.pdf"
        response = requests.get(url)
        response.raise_for_status()
        pdf_content = BytesIO(response.content)
        reader = PdfReader(pdf_content)
        text = ""
        for page in reader.pages:
            text += page.extract_text()
        return text
    except Exception as e:
        print(f"Error fetching paper content: {e}")
        return None


def fetch_paper_info_paperpage(paper_id_value):
    # Extract paper_id from paper_page link or input
    def extract_paper_id(input_string):
        # Already in correct form?
        if re.fullmatch(r'\d+\.\d+', input_string.strip()):
            return input_string.strip()
        # If URL
        match = re.search(r'https://huggingface\.co/papers/(\d+\.\d+)', input_string)
        if match:
            return match.group(1)
        return input_string.strip()

    paper_id_value = extract_paper_id(paper_id_value)
    url = f"https://huggingface.co/api/papers/{paper_id_value}?field=comments"
    response = requests.get(url)
    if response.status_code != 200:
        return None, None, None
    paper_info = response.json()
    title = paper_info.get('title', 'No Title')
    authors_list = [author.get('name', 'Unknown') for author in paper_info.get('authors', [])]
    authors = ', '.join(authors_list)
    summary = paper_info.get('summary', 'No Summary')
    num_comments = len(paper_info.get('comments', []))
    num_upvotes = paper_info.get('upvotes', 0)
    link = f"https://huggingface.co/papers/{paper_id_value}"

    details = f"{summary}<br/>👍{num_comments} 💬{num_upvotes}<br/> <a href='{link}' " \
              f"target='_blank'>View on 🤗 hugging face</a>"
    return title, authors, details


def fetch_paper_content_paperpage(paper_id_value):
    # Extract paper_id
    def extract_paper_id(input_string):
        if re.fullmatch(r'\d+\.\d+', input_string.strip()):
            return input_string.strip()
        match = re.search(r'https://huggingface\.co/papers/(\d+\.\d+)', input_string)
        if match:
            return match.group(1)
        return input_string.strip()

    paper_id_value = extract_paper_id(paper_id_value)
    text = fetch_paper_content_arxiv(paper_id_value)
    return text


# Dictionary for paper sources
PAPER_SOURCES = {
    "neurips": {
        "fetch_info": fetch_paper_info_neurips,
        "fetch_pdf": fetch_paper_content_neurips
    },
    "paper_page": {
        "fetch_info": fetch_paper_info_paperpage,
        "fetch_pdf": fetch_paper_content_paperpage
    }
}


def create_chat_interface(provider_dropdown, model_dropdown, paper_content, hf_token_input, default_type,
                          provider_max_total_tokens):
    # Define the function to handle the chat
    def get_fn(message, history, paper_content_value, hf_token_value, provider_name_value, model_name_value,
               max_total_tokens):
        provider_info = PROVIDERS[provider_name_value]
        endpoint = provider_info['endpoint']
        api_key_env_var = provider_info['api_key_env_var']
        models = provider_info['models']
        max_total_tokens = int(max_total_tokens)

        # Load tokenizer
        tokenizer_key = f"{provider_name_value}_{model_name_value}"
        if tokenizer_key not in tokenizer_cache:
            tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B-Instruct",
                                                      token=os.environ.get("HF_TOKEN"))
            tokenizer_cache[tokenizer_key] = tokenizer
        else:
            tokenizer = tokenizer_cache[tokenizer_key]

        # Include the paper content as context
        if paper_content_value:
            context = f"The discussion is about the following paper:\n{paper_content_value}\n\n"
        else:
            context = ""

        # Tokenize the context
        context_tokens = tokenizer.encode(context)
        context_token_length = len(context_tokens)

        # Prepare the messages without context
        messages = []
        message_tokens_list = []
        total_tokens = context_token_length  # Start with context tokens

        for user_msg, assistant_msg in history:
            # Tokenize user message
            user_tokens = tokenizer.encode(user_msg)
            messages.append({"role": "user", "content": user_msg})
            message_tokens_list.append(len(user_tokens))
            total_tokens += len(user_tokens)

            # Tokenize assistant message
            if assistant_msg:
                assistant_tokens = tokenizer.encode(assistant_msg)
                messages.append({"role": "assistant", "content": assistant_msg})
                message_tokens_list.append(len(assistant_tokens))
                total_tokens += len(assistant_tokens)

        # Tokenize the new user message
        message_tokens = tokenizer.encode(message)
        messages.append({"role": "user", "content": message})
        message_tokens_list.append(len(message_tokens))
        total_tokens += len(message_tokens)

        # Check if total tokens exceed the maximum allowed tokens
        if total_tokens > max_total_tokens:
            # Attempt to truncate context
            available_tokens = max_total_tokens - (total_tokens - context_token_length)
            if available_tokens > 0:
                truncated_context_tokens = context_tokens[:available_tokens]
                context = tokenizer.decode(truncated_context_tokens)
                context_token_length = available_tokens
                total_tokens = total_tokens - len(context_tokens) + context_token_length
            else:
                context = ""
                total_tokens -= context_token_length
                context_token_length = 0

        # Truncate message history if needed
        while total_tokens > max_total_tokens and len(messages) > 1:
            removed_message = messages.pop(0)
            removed_tokens = message_tokens_list.pop(0)
            total_tokens -= removed_tokens

        # Rebuild the final messages
        final_messages = []
        if context:
            final_messages.append({"role": "system", "content": f"{context}"})
        final_messages.extend(messages)

        # Use the provider's API key
        api_key = hf_token_value or os.environ.get(api_key_env_var)
        if not api_key:
            raise ValueError("API token is not provided.")

        # Initialize the OpenAI client
        client = OpenAI(
            base_url=endpoint,
            api_key=api_key,
        )

        try:
            completion = client.chat.completions.create(
                model=model_name_value,
                messages=final_messages,
                stream=True,
            )
            response_text = ""
            for chunk in completion:
                delta = chunk.choices[0].delta.content or ""
                response_text += delta
                yield response_text
        except json.JSONDecodeError as e:
            yield f"JSON decoding error: {e.msg}"
        except openai.OpenAIError as openai_err:
            yield f"OpenAI error: {openai_err}"
        except Exception as ex:
            yield f"Unexpected error: {ex}"

    chatbot = gr.Chatbot(label="Chatbot", scale=1, height=400, autoscroll=True)
    chat_interface = gr.ChatInterface(
        fn=get_fn,
        chatbot=chatbot,
        additional_inputs=[paper_content, hf_token_input, provider_dropdown, model_dropdown, provider_max_total_tokens],
        type="tuples",
    )
    return chat_interface, chatbot


def paper_chat_tab(paper_id, paper_from, paper_central_df):
    with gr.Row():
        # Left column: Paper selection and display
        with gr.Column(scale=1):
            gr.Markdown("### Select a Paper")
            todays_date = datetime.today().strftime('%Y-%m-%d')

            # Filter papers for today's date and having a paper_page
            selectable_papers = paper_central_df.df_prettified
            selectable_papers = selectable_papers[
                selectable_papers['paper_page'].notna() &
                (selectable_papers['paper_page'] != "") &
                (selectable_papers['date'] == todays_date)
                ]

            paper_choices = [(row['title'], row['paper_page']) for _, row in selectable_papers.iterrows()]
            paper_choices = sorted(paper_choices, key=lambda x: x[0])

            if not paper_choices:
                paper_choices = [("No available papers for today", "")]

            paper_select = gr.Dropdown(
                label="Select a paper to chat with:",
                choices=[p[0] for p in paper_choices],
                value=paper_choices[0][0] if paper_choices else None
            )
            select_paper_button = gr.Button("Load this paper")

            # Paper info display - styled card
            content = gr.HTML(value="", elem_id="paper_info_card")

        # Right column: Provider and model selection + chat
        with gr.Column(scale=1, visible=False) as provider_section:
            gr.Markdown("### LLM Provider and Model")
            provider_names = list(PROVIDERS.keys())
            default_provider = provider_names[0]

            default_type = gr.State(value=PROVIDERS[default_provider]["type"])
            default_max_total_tokens = gr.State(value=PROVIDERS[default_provider]["max_total_tokens"])

            provider_dropdown = gr.Dropdown(
                label="Select Provider",
                choices=provider_names,
                value=default_provider
            )

            hf_token_input = gr.Textbox(
                label=f"Enter your {default_provider} API token (optional)",
                type="password",
                placeholder=f"Enter your {default_provider} API token to avoid rate limits"
            )

            model_dropdown = gr.Dropdown(
                label="Select Model",
                choices=PROVIDERS[default_provider]['models'],
                value=PROVIDERS[default_provider]['models'][0]
            )

            logo_html = gr.HTML(
                value=f'<img src="{PROVIDERS[default_provider]["logo"]}" width="100px" />'
            )

            note_markdown = gr.Markdown(f"**Note:** This model is supported by {default_provider}.")

            paper_content = gr.State()

            # Create chat interface
            chat_interface, chatbot = create_chat_interface(provider_dropdown, model_dropdown, paper_content,
                                                            hf_token_input, default_type, default_max_total_tokens)

    def update_provider(selected_provider):
        provider_info = PROVIDERS[selected_provider]
        models = provider_info['models']
        logo_url = provider_info['logo']
        chatbot_message_type = provider_info['type']
        max_total_tokens = provider_info['max_total_tokens']

        model_dropdown_choices = gr.update(choices=models, value=models[0])
        logo_html_content = f'<img src="{logo_url}" width="100px" />'
        logo_html_update = gr.update(value=logo_html_content)
        note_markdown_update = gr.update(value=f"**Note:** This model is supported by {selected_provider}.")
        hf_token_input_update = gr.update(
            label=f"Enter your {selected_provider} API token (optional)",
            placeholder=f"Enter your {selected_provider} API token to avoid rate limits"
        )
        chatbot_reset = []
        return model_dropdown_choices, logo_html_update, note_markdown_update, hf_token_input_update, chatbot_message_type, max_total_tokens, chatbot_reset

    provider_dropdown.change(
        fn=update_provider,
        inputs=provider_dropdown,
        outputs=[model_dropdown, logo_html, note_markdown, hf_token_input, default_type, default_max_total_tokens,
                 chatbot],
        queue=False
    )

    def update_paper_info(paper_id_value, paper_from_value, selected_model, old_content):
        # Use PAPER_SOURCES to fetch info
        source_info = PAPER_SOURCES.get(paper_from_value, {})
        fetch_info_fn = source_info.get("fetch_info")
        fetch_pdf_fn = source_info.get("fetch_pdf")

        if not fetch_info_fn or not fetch_pdf_fn:
            return gr.update(value="<div>No information available.</div>"), None, []

        title, authors, details = fetch_info_fn(paper_id_value)
        if title is None and authors is None and details is None:
            return gr.update(value="<div>No information could be retrieved.</div>"), None, []

        text = fetch_pdf_fn(paper_id_value)
        if text is None:
            text = "Paper content could not be retrieved."

        # Create a styled card for the paper info
        card_html = f"""
        <div style="border:1px solid #ccc; border-radius:6px; background:#f9f9f9; padding:15px; margin-bottom:10px;">
            <center><h3 style="margin-top:0; text-decoration:underline;">You are talking with:</h3></center>
            <h3>{title}</h3>
            <p><strong>Authors:</strong> {authors}</p>
            <p>{details}</p>
        </div>
        """

        return gr.update(value=card_html), text, []

    def select_paper(paper_title):
        # Find the corresponding paper_page from the title
        for t, ppage in paper_choices:
            if t == paper_title:
                return ppage, "paper_page"
        return "", ""

    select_paper_button.click(
        fn=select_paper,
        inputs=[paper_select],
        outputs=[paper_id, paper_from]
    )

    # After updating paper_id, we update paper info
    paper_id.change(
        fn=update_paper_info,
        inputs=[paper_id, paper_from, model_dropdown, content],
        outputs=[content, paper_content, chatbot]
    )

    # Function to toggle visibility of the right column based on paper_id
    def toggle_provider_visibility(paper_id_value):
        if paper_id_value and paper_id_value.strip():
            return gr.update(visible=True)
        else:
            return gr.update(visible=False)

    # Chain a then call to toggle visibility of the provider_section after paper info update
    paper_id.change(
        fn=toggle_provider_visibility,
        inputs=[paper_id],
        outputs=[provider_section]
    )


def main():
    """
    Launches the Gradio app.
    """
    with gr.Blocks(css_paths="style.css") as demo:
        paper_id = gr.Textbox(label="Paper ID", value="")
        paper_from = gr.Radio(
            label="Paper Source",
            choices=["neurips", "paper_page"],
            value="neurips"
        )

        # Build the paper chat tab
        dummy_calendar = gr.State(datetime.now().strftime("%Y-%m-%d"))

        class MockPaperCentral:
            def __init__(self):
                import pandas as pd
                data = {
                    'date': [datetime.today().strftime('%Y-%m-%d')],
                    'paper_page': ['1234.56789'],
                    'title': ['An Example Paper']
                }
                self.df_prettified = pd.DataFrame(data)

        paper_central_df = MockPaperCentral()

        paper_chat_tab(paper_id, paper_from, paper_central_df)

    demo.launch(ssr_mode=False)


if __name__ == "__main__":
    main()