import streamlit as st
from PIL import Image
import os
from datetime import datetime
from hf_model import HFModel  # Import your HFModel class
from mew_log import log_info, log_error  # Import your custom logging methods

# Initialize the model
MODEL_NAME = "deepseek-ai/DeepSeek-Coder-V2-Lite-Base"

def initialize_model():
    try:
        model = HFModel(MODEL_NAME)
        log_info("Model initialized successfully.")
        return model
    except Exception as e:
        log_error(f"Failed to initialize model: {e}", e)
        st.error(f"Failed to initialize model: {e}")
        return None

# Initialize the model (cached to avoid reloading on every interaction)
@st.cache_resource
def load_model():
    return initialize_model()

# Streamlit app
def main():
    st.title("DeepSeek Chat App")
    st.write("Interact with the DeepSeek model using text, images, or files!")

    # Load the model
    model = load_model()
    if model is None:
        st.stop()

    # Sidebar for file uploads
    st.sidebar.header("Upload Files or Images")
    uploaded_file = st.sidebar.file_uploader("Upload a file or image", type=["txt", "jpg", "png", "jpeg", "pdf"])

    # Display uploaded file or image
    if uploaded_file is not None:
        file_details = {"filename": uploaded_file.name, "filetype": uploaded_file.type, "filesize": uploaded_file.size}
        st.sidebar.write(file_details)

        if uploaded_file.type.startswith("image"):
            image = Image.open(uploaded_file)
            st.sidebar.image(image, caption="Uploaded Image", use_column_width=True)
        else:
            st.sidebar.write("File content:")
            file_content = uploaded_file.getvalue().decode("utf-8")
            st.sidebar.text(file_content)

    # Main chat interface
    st.header("Chat with DeepSeek")
    user_input = st.text_area("Enter your message or question:", height=100)

    if st.button("Send"):
        if user_input.strip() == "":
            st.warning("Please enter a message.")
        else:
            try:
                # Add file content to user input if a file is uploaded
                if uploaded_file is not None:
                    if uploaded_file.type.startswith("image"):
                        user_input += f"\n[Attached Image: {uploaded_file.name}]"
                    else:
                        file_content = uploaded_file.getvalue().decode("utf-8")
                        user_input += f"\n[Attached File Content:\n{file_content}\n]"

                # Generate response
                with st.spinner("Generating response..."):
                    response = model.chat(user_input, max_length=512)
                    st.write(f"**Assistant:** {response}")

                # Display chat history
                st.subheader("Chat History")
                for entry in model.chat_history:
                    role = entry["role"]
                    content = entry["content"]
                    st.write(f"**{role.capitalize()}:** {content}")

            except Exception as e:
                log_error(f"Error during chat: {e}", e)
                st.error(f"An error occurred: {e}")

    # Clear chat history button
    if st.button("Clear Chat History"):
        model.clear_chat_history()
        st.success("Chat history cleared!")

# Run the app
if __name__ == "__main__":
    main()