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 = ["Rosenshine+Principles.pdf", AnotherDocument.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.
    Use the following as context for your answers.
    {context}
    However, use it seamlessly as background knowledge for a lively discussion. 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.
    """

    # Customize instructions as needed
    instructions = instructions
    system_message += instructions

    # 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=500,
        system=system_message,
        messages=messages
    )

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

# CSS for an ISP-looking style
isp_theme = gr.themes.Default().set(
    body_background_fill="#F0F8FF",  # Light blue background
    block_background_fill="#FFFFFF",  # White for input blocks
    block_title_text_color="#00008B",  # Dark blue for text
    block_label_background_fill="#FFD700",  # Yellow for labels
    input_background_fill="#FFFFFF",  # White for input fields
    button_primary_background_fill="#FF1493",  # Pink/Magenta for primary buttons
    button_primary_background_fill_hover="#FF69B4",  # Lighter pink for hover
    button_primary_text_color="#FFFFFF",  # White text on buttons
    button_secondary_background_fill="#FFD700",  # Yellow for secondary buttons
    button_secondary_background_fill_hover="#FFA500",  # Orange for hover
    button_secondary_text_color="#00008B",  # Dark blue text for secondary buttons
    block_border_width="1px",
    block_border_color="#00008B",  # Dark blue border
)

# Custom CSS for logo positioning
custom_css = """
#logo-img {
    position: absolute;
    top: 20px;
    right: 20px;
    width: 100px;
    height: auto;
}
"""

assistant_avatar = os.getenv('AVATAR')
assistant_title = os.getenv('TITLE')
# Gradio interface
iface = gr.ChatInterface(
    chat_with_assistant,
    chatbot=gr.Chatbot(
        height=500, 
        avatar_images=(None, assistant_avatar)
    ),
    textbox=gr.Textbox(placeholder="Type your message here...", container=False, scale=7),
    title=assistant_title,
    description="Chat with a Historical Figure",
    theme=isp_theme,
    retry_btn=None,
    undo_btn="Delete Previous",
    clear_btn="Clear",
    css=custom_css,
)

assistant_logo = os.getenv('LOGO')

# Add logo to the interface
logo_html = '&lt;img id="logo-img" src=assistant_logo alt="Assistant Logo"&gt;'
iface.attachments.append(logo_html)

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