import streamlit as st
import faiss
import langchain
from llama_index.callbacks import CallbackManager
from llama_index import  ServiceContext,VectorStoreIndex
from llama_index.chat_engine import CondenseQuestionChatEngine
from llama_index.node_parser import SimpleNodeParser
from llama_index.langchain_helpers.text_splitter import TokenTextSplitter
from llama_index.constants import DEFAULT_CHUNK_OVERLAP
from llama_index.response_synthesizers import get_response_synthesizer
from llama_index import  SimpleWebPageReader

# from llama_index.prompts import Prompt
from llama_index import Prompt
import tiktoken
import common
langchain.verbose = True

custom_prompt = Prompt("""\
  以下はこれまでの会話履歴と、ドキュメントを検索して回答する必要がある、ユーザーからの会話文です。
  会話と新しい会話文に基づいて、検索クエリを作成します。回答は日本語で行います。
  新しい会話文が挨拶の場合、挨拶を返してください。
  新しい会話文が質問の場合、検索した結果の回答を返してください。
  答えがわからない場合は正直にわからないと回答してください。
  会話履歴:
  {chat_history}
  新しい会話文:
  {question}
  Search query:
""")
chat_history = []

from log import logger

common.check_login()

st.title("💬 ChatbotWebRead")

URLtext = st.text_input(
    "読み込むURLを入力してください",
    placeholder="https://",
)

if st.button("URL reading",use_container_width=True):
    text_splitter = TokenTextSplitter( chunk_size=1500
      , chunk_overlap=DEFAULT_CHUNK_OVERLAP
      , tokenizer=tiktoken.encoding_for_model("gpt-3.5-turbo").encode)
    node_parser = SimpleNodeParser(text_splitter=text_splitter)
    d = 1536
    k=2
    faiss_index = faiss.IndexFlatL2(d)

    callback_manager = CallbackManager([st.session_state.llama_debug_handler])
    service_context = ServiceContext.from_defaults(node_parser=node_parser,callback_manager=callback_manager)

    webDocuments = SimpleWebPageReader(html_to_text=True).load_data(
        [URLtext]
    )
    logger.info(webDocuments)
    webIndex = VectorStoreIndex.from_documents(webDocuments,service_context=service_context)
    response_synthesizer = get_response_synthesizer(response_mode='refine')
    st.session_state.webQuery_engine = webIndex.as_query_engine(
        response_synthesizer=response_synthesizer,
        service_context=service_context,
    )
    st.session_state.web_chat_engine = CondenseQuestionChatEngine.from_defaults(
        query_engine=st.session_state.webQuery_engine, 
        condense_question_prompt=custom_prompt,
        chat_history=chat_history,
        verbose=True
    )

if st.button("リセット",use_container_width=True,disabled = not URLtext):
    st.session_state.web_chat_engine.reset()
    st.session_state.webmessages = [{"role": "assistant", "content": "お困りごとはございますか？"}]
    st.experimental_rerun()
    logger.info("reset")

if "webmessages" not in st.session_state:
    st.session_state["webmessages"] = [{"role": "assistant", "content": "お困りごとはございますか？"}]

for msg in st.session_state.webmessages:
    st.chat_message(msg["role"]).write(msg["content"])

if prompt := st.chat_input(disabled = not URLtext):
    st.session_state.webmessages.append({"role": "user", "content": prompt})
    st.chat_message("user").write(prompt)
    response = st.session_state.web_chat_engine.chat(prompt)
    logger.debug(st.session_state.llama_debug_handler.get_llm_inputs_outputs())
    msg = str(response)
    st.session_state.webmessages.append({"role": "assistant", "content": msg})
    st.chat_message("assistant").write(msg)