# Combined Imports
import os
import streamlit as st
from dotenv import load_dotenv
from apify_client import ApifyClient
from langchain.callbacks.base import BaseCallbackHandler
from langchain.chains import ConversationalRetrievalChain
from langchain.chat_models import ChatOpenAI
from langchain.document_loaders import ApifyDatasetLoader
from langchain.document_loaders.base import Document
from langchain.embeddings import OpenAIEmbeddings
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.memory import ConversationBufferMemory
from langchain.memory.chat_message_histories import StreamlitChatMessageHistory
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma

# Environment variables and configuration
load_dotenv()
WEBSITE_URL = os.environ.get('WEBSITE_URL', 'a website')
OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
APIFY_API_TOKEN = os.environ.get('APIFY_API_TOKEN')

# Scraper Functionality
def scrape_website():
    apify_client = ApifyClient(APIFY_API_TOKEN)
    st.write(f'Extracting data from "{WEBSITE_URL}". Please wait...')
    actor_run_info = apify_client.actor('apify/website-content-crawler').call(
        run_input={'startUrls': [{'url': WEBSITE_URL}]}
    )
    st.write('Saving data into the vector database. Please wait...')
    loader = ApifyDatasetLoader(
        dataset_id=actor_run_info['defaultDatasetId'],
        dataset_mapping_function=lambda item: Document(
            page_content=item['text'] or '', metadata={'source': item['url']}
        ),
    )
    documents = loader.load()
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1500, chunk_overlap=100)
    docs = text_splitter.split_documents(documents)

    embedding = OpenAIEmbeddings(api_key=OPENAI_API_KEY)
    vectordb = Chroma.from_documents(
        documents=docs,
        embedding=embedding,
        persist_directory='db2',
    )
    vectordb.persist()
    st.write('All done!')

# Chat Functionality
def chat_with_website():
    st.set_page_config(page_title=f'Chat with {WEBSITE_URL}')
    st.title('Chat with a website')
    retriever = get_retriever()
    msgs = StreamlitChatMessageHistory()
    memory = ConversationBufferMemory(memory_key='chat_history', chat_memory=msgs, return_messages=True)

    llm = ChatOpenAI(model_name='gpt-3.5-turbo', temperature=0, streaming=True)
    qa_chain = ConversationalRetrievalChain.from_llm(
        llm, retriever=retriever, memory=memory, verbose=False
    )

    if st.sidebar.button('Clear message history') or len(msgs.messages) == 0:
        msgs.clear()
        msgs.add_ai_message(f'Ask me anything about {WEBSITE_URL}!')

    avatars = {'human': 'user', 'ai': 'assistant'}
    for msg in msgs.messages:
        st.chat_message(avatars[msg.type]).write(msg.content)

    if user_query := st.chat_input(placeholder='Ask me anything!'):
        st.chat_message('user').write(user_query)
        with st.chat_message('assistant'):
            stream_handler = StreamHandler(st.empty())
            response = qa_chain.run(user_query, callbacks=[stream_handler])

@st.cache_resource(ttl='1h')
def get_retriever():
    embeddings = OpenAIEmbeddings()
    vectordb = Chroma(persist_directory='db', embedding_function=embeddings)
    retriever = vectordb.as_retriever(search_type='mmr')
    return retriever

class StreamHandler(BaseCallbackHandler):
    def __init__(self, container: st.delta_generator.DeltaGenerator, initial_text: str = ''):
        self.container = container
        self.text = initial_text

    def on_llm_new_token(self, token: str, **kwargs) -> None:
        self.text += token
        self.container.markdown(self.text)

# Main App Flow
if st.sidebar.button("Scrape a new website"):
    scrape_website()

if st.sidebar.button("Chat with scraped website"):
    chat_with_website()