import streamlit as st
from transformers import pipeline
import os
import PyPDF2

# Sidebar components
st.sidebar.title("TherapyGPT Settings")

# Model selection dropdown
model_option = st.sidebar.selectbox(
    "Choose AI Model:",
    ("GPT-4o", "distilgpt2", "gpt-neo")
)

# Document upload
uploaded_file = st.sidebar.file_uploader(
    "Upload Document", type=["txt", "pdf"])

# Workspace/session management
workspace_name = st.sidebar.text_input("Workspace Name", value="Default")
if st.sidebar.button("Create Workspace"):
    st.session_state.workspace = workspace_name
    st.session_state.chat_history = []  # Reset chat history


# Main chat area
st.title("TherapyGPT Chat")

# Display chat history if it exists
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []

for message in st.session_state.chat_history:
    st.write(f"{message['role']}: {message['content']}")

# User input
user_input = st.text_input("You:", "")
if st.button("Send"):
    if user_input:
        # Placeholder for model interaction (we'll integrate this next)
        st.session_state.chat_history.append(
            {"role": "User", "content": user_input})
        response = "I'm here to help."  # Temporary placeholder
        st.session_state.chat_history.append(
            {"role": "TherapyGPT", "content": response})


# Load the model based on selection

@st.cache_resource
def load_model(model_name):
    return pipeline("text-generation", model=model_name)


# Model selection logic
if model_option == "GPT-4o":
    # Replace with appropriate model
    model = load_model("ruslandev/llama-3-8b-gpt-4o-ru1.0-gguf")
elif model_option == "distilgpt2":
    pipe = pipeline("text-generation", model="distilbert/distilgpt2")
else:
    # Example HuggingFace model
    pipe = pipeline("text-generation", model="EleutherAI/gpt-neo-1.3B")


def read_file(file):
    if file.name.endswith('.txt'):
        return file.read().decode('utf-8')
    elif file.name.endswith('.pdf'):
        pdf_reader = PyPDF2.PdfReader(file)
        return '\n'.join([page.extract_text() for page in pdf_reader.pages])


if uploaded_file:
    file_content = read_file(uploaded_file)
    st.sidebar.text_area("File Content", file_content, height=200)


if 'workspaces' not in st.session_state:
    st.session_state.workspaces = {}

# Load or create the workspace
if workspace_name not in st.session_state.workspaces:
    st.session_state.workspaces[workspace_name] = []

# Use the chat history for this workspace
chat_history = st.session_state.workspaces[workspace_name]