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 = ["Primary sources for Summative Project.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 speak in the first person ("I") as the historical figure you are to incarnate.
    
    Always use appropriate language.
    Refuse to answer inappropriate questions or questions unrelated to your role and historical figure.
    Keep your responses concise and to the point. Avoid repetitions and always end on a period.

    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=250,
        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))