import os
import gradio as gr
from anthropic import Anthropic
from pypdf import PdfReader
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

# Set up Anthropic API key in HF secrets
ANTHROPIC_API_KEY = os.getenv('ANTHROPIC_API_KEY')
os.environ["ANTHROPIC_API_KEY"] = ANTHROPIC_API_KEY

# Set up username and password in HF secrets
username = os.getenv('username')
password = os.getenv('password')

# Function to chunk the document
def chunk_text(text, chunk_size=1000, overlap=100):
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunk = text[start:end]
        chunks.append(chunk)
        start = end - overlap
    return chunks

# Function to find the most relevant chunks
def get_relevant_chunks(query, chunks, top_n=3):
    vectorizer = TfidfVectorizer()
    tfidf_matrix = vectorizer.fit_transform(chunks + [query])
    cosine_similarities = cosine_similarity(tfidf_matrix[-1], tfidf_matrix[:-1]).flatten()
    relevant_indices = cosine_similarities.argsort()[-top_n:][::-1]
    return [chunks[i] for i in relevant_indices]

# Function to process multiple PDFs
def process_pdfs(pdf_files):
    all_chunks = []
    for pdf_file in pdf_files:
        reader = PdfReader(pdf_file)
        full_text = ''.join(page.extract_text() for page in reader.pages)
        chunks = chunk_text(full_text)
        all_chunks.extend(chunks)
    return all_chunks

# Add the paths to your desired knowledge base PDFs
reference_documents = ["Louis XIV.pdf"]
text_chunks = process_pdfs(reference_documents)

instructions = os.getenv('INSTRUCTIONS')

def chat_with_assistant(message, history):
    # Find relevant chunks based on the user message
    relevant_chunks = get_relevant_chunks(message, text_chunks)
    context = "\n".join(relevant_chunks)

    # Prepare the system message
    system_message = f"""You are an impersonator and an educator.
    Your role is to adopt the personality, style, psychology, ideas, background, and circumstances of a historical figure.
    Your goal is to help students understand the historical figure better through and engaging conversation.

    Your assigned historical figure is stated in your instructions: 
    {instructions}
    
    Use the following as context for your answers.
    {context}
    However, use it seamlessly as background knowledge for a lively discussion and combine it with your own information. Do not provide citations or adopt a Q&A or academic tone.
    
    Always use appropriate language.
    Refuse to answer inappropriate questions or questions unrelated to your role and historical figure.

    Important: Your knowledge of the world ends at the time of the death of your historical figure.
    """

    # Prepare the message array
    messages = []

    # Add conversation history
    for human_msg, ai_msg in history:
        messages.append({"role": "user", "content": human_msg})
        messages.append({"role": "assistant", "content": ai_msg})

    # Add the current user message
    messages.append({"role": "user", "content": message})

    # Create Anthropic client
    client = Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])

    # Make the API call
    response = client.messages.create(
        model="claude-3-5-sonnet-20240620",
        max_tokens=350,
        system=system_message,
        messages=messages
    )

    return response.content[0].text.strip()

# CSS for a blue-themed style
isp_theme = gr.themes.Default().set(
    body_background_fill="#E6F3FF",  # Light blue background
    block_background_fill="#FFFFFF",  # White for input blocks
    block_title_text_color="#003366",  # Dark blue for text
    block_label_background_fill="#B8D8FF",  # Light blue for labels
    input_background_fill="#FFFFFF",  # White for input fields
    button_primary_background_fill="#0066CC",  # Medium blue for primary buttons
    button_primary_background_fill_hover="#0052A3",  # Darker blue for hover
    button_primary_text_color="#FFFFFF",  # White text on buttons
    button_secondary_background_fill="#B8D8FF",  # Light blue for secondary buttons
    button_secondary_background_fill_hover="#99C2FF",  # Slightly darker blue for hover
    button_secondary_text_color="#003366",  # Dark blue text for secondary buttons
    block_border_width="1px",
    block_border_color="#0066CC",  # Medium blue border
)

# Custom CSS for logo positioning and disclaimer footer
custom_css = """
#logo-img { 
    display: block;
    margin: 0 auto;
    width: 150px; 
    height: auto;
    padding-bottom: 20px; /* Space below logo */
}
#disclaimer-footer { 
    width: 100%; 
    background-color: #B8D8FF; 
    color: #003366; 
    text-align: center; 
    padding: 10px 0; 
    font-size: 14px; 
    border-top: 1px solid #0066CC; 
    margin-top: 20px;
}
.container { 
    max-width: 1200px; 
    margin: 0 auto; 
    padding: 10px;
}
.title { 
    color: #003366; 
    margin-bottom: 10px;
    text-align: center; /* Center the title */
}
.chatbot { 
    border: none;  
    border-radius: 5px; 
    padding: 10px;  
    margin-bottom: 15px;
}
.button-row { 
    display: flex; 
    gap: 10px; 
    justify-content: center; /* Center the buttons */
    margin-bottom: 15px;
}
.chatbot .message,
.chatbot .message::before,
.chatbot .message::after {
    border: none !important;
    box-shadow: none !important;
}
.chatbot .message > div {
    border: none !important;
    box-shadow: none !important;
}
.chatbot .message-content {
    padding: 2px 2px; /* Reduced padding to make text bubbles smaller */
    margin-bottom: 5px; /* Space between messages */
    max-width: 60%; /* Restrict width to make it more compact */
    word-wrap: break-word; /* Ensure text wraps properly */
    display: inline-block; /* Align content properly */
}
.chatbot .message-bubble {
    background-color: #FFFFFF; /* Ensure background is white */
    border-radius: 2px; /* Smaller rounded corners */
    box-sizing: border-box; /* Ensure padding is included in width */
    display: inline-block; /* Align the bubble content properly */
    margin: 2px 0; /* Reduce margin to minimize bubble size */
}
"""

# Environment variables
assistant_avatar = os.getenv('AVATAR')
assistant_title = os.getenv('TITLE')
assistant_logo = os.getenv('LOGO')

# Gradio interface using Blocks
with gr.Blocks(theme=isp_theme, css=custom_css) as iface:
    with gr.Column(elem_classes="container"):
        # Logo and Title
        gr.HTML(
            f"""
            <img id='logo-img' src='{assistant_logo}' alt='Assistant Logo' onerror="this.style.display='none';document.getElementById('logo-error').style.display='block';">
            <div id='logo-error' style='display:none;'>Logo not found</div>
            """,
            elem_id='logo-container'
        )
        gr.Markdown(f"# {assistant_title}", elem_classes="title")

        # Chatbot and message input
        with gr.Row():
            chatbot = gr.Chatbot(
                height=500,
                avatar_images=(None, assistant_avatar),
                elem_classes="chatbot"
            )

        msg = gr.Textbox(
            placeholder="Type your message here...",
            container=False,
            scale=7
        )

        with gr.Row(elem_classes="button-row"):
            submit = gr.Button("Submit", variant="primary")
            clear = gr.ClearButton([msg, chatbot], value="Clear", variant="secondary")
            undo = gr.Button("Delete Previous", variant="secondary")

        gr.HTML(
            "<div id='disclaimer-footer'>You are chatting with an AI assistant. Make sure to evaluate the accuracy of its answers.</div>"
        )

    def user(user_message, history):
        return "", history + [[user_message, None]]

    def bot(history):
        bot_message = chat_with_assistant(history[-1][0], history[:-1])
        history[-1][1] = bot_message
        return history

    def delete_previous(history):
        if len(history) > 0:
            return history[:-1]
        return history

    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, chatbot, chatbot
    )
    submit.click(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, chatbot, chatbot
    )
    undo.click(delete_previous, chatbot, chatbot)

iface.launch(auth=(username, password))