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 { 
    position: absolute; 
    top: 20px; 
    right: 20px; 
    width: 200px; 
    height: auto;
    z-index: 1000;
}
#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: 800px; 
    margin: 0 auto; 
    padding: 20px;
    position: relative;
    padding-top: 60px;
}
.title { 
    color: #003366; 
    margin-bottom: 10px;
}
.chatbot { 
    border: none;  /* Remove the blue border */
    border-radius: 5px; 
    padding: 5px;  /* Reduce padding around the chat area */
    margin-bottom: 15px;
}
.button-row { 
    display: flex; 
    gap: 10px; 
    margin-bottom: 15px;
}
/* Remove borders and shadows from chat bubbles */
.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;
}
/* Adjust padding, margin, and width for chat bubbles */
.chatbot .message-content {
    padding: 5px 10px; /* Smaller padding for better sizing */
    margin-bottom: 8px; /* Space between messages */
    max-width: 70%; /* Restrict maximum width for better appearance */
    word-wrap: break-word; /* Ensure text wraps properly */
}
.chatbot .message-bubble {
    background-color: #FFFFFF; /* Ensure background is white */
    border-radius: 8px; /* Rounded corners for messages */
    box-sizing: border-box; /* Ensure padding is included in width */
    display: inline-block; /* Align the bubble content properly */
}
"""

# 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"):
        with gr.Row():
            # 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")

        gr.Markdown("Chat with a Historical Figure")

        # Chatbot and message input
        with gr.Row():
            with gr.Column(scale=6):
                chatbot = gr.Chatbot(
                    height=500,
                    avatar_images=(None, assistant_avatar),
                    elem_classes="chatbot"
                )
            with gr.Column(scale=1):
                msg = gr.Textbox(
                    placeholder="Type your message here...",
                    container=False
                )
                
        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))