import json
import os
import re

import chainlit as cl
from dotenv import load_dotenv
from langchain.embeddings import CacheBackedEmbeddings
from langchain.prompts import ChatPromptTemplate
from langchain.schema import StrOutputParser
from langchain.schema.runnable import RunnablePassthrough
from langchain.schema.runnable.config import RunnableConfig
from langchain.storage import LocalFileStore
from langchain_core.runnables import RunnableParallel
from langchain_openai import ChatOpenAI
from langchain_openai import OpenAIEmbeddings
from langchain_pinecone import PineconeVectorStore
import requests

load_dotenv()

RAG_PROMPT = """
CONTEXT:
{context}
QUERY:
{question}

You are a helpful physical therapy assistant. You provide summarized
and succint information. Your answers are accurate and thorough.
You will be presented with a question helping
one of our physical therapist clients, and your job is to be as detailed
and as helpful as possible. If the context provided
doesn't answer the question, please respond with: "I
require more information in order to better assist you,
please state your question and what kind of service or
support you are seeking."
"""

core_embeddings_model = OpenAIEmbeddings()
vector_store = PineconeVectorStore(index_name=os.getenv(
    "INDEX_NAME"), embedding=core_embeddings_model)
retriever = vector_store.as_retriever(search_kwargs={"k": 1})

llm = ChatOpenAI(streaming=True, model=os.getenv("OPENAI_MODEL"))
rag_prompt = ChatPromptTemplate.from_template(RAG_PROMPT)


def format_docs(docs):
    return "\n\n".join(doc for doc in docs)


rag_chain_from_docs = (
    RunnablePassthrough.assign(context=(lambda x: format_docs(x)))
    | rag_prompt
    | llm
    | StrOutputParser()
)

rag_chain_with_source = RunnableParallel(
    {"context": retriever, "question": RunnablePassthrough()}
).assign(answer=rag_chain_from_docs)


@cl.on_chat_start
async def on_chat_start():
    cl.user_session.set("runnable", rag_chain_with_source)


@cl.on_message
async def on_message(message: cl.Message):
    rag_chain_with_source = cl.user_session.get("runnable")
    msg = cl.Message(content="")
    response = rag_chain_with_source.invoke(message.content)

    if response is not None:
        metadata = response["context"][0].metadata
        link = metadata["link"]
        title = metadata["source_document"]

        print(response['answer'])

        formatted_response = f"""
            {response['answer']} Check out the Youtube video [{title}] ({link})!
        """

        match = re.search(r"<(\d+\.\d+)>", response["context"][0].page_content)
        if match:
            start_time = float(match.group(1))
            formatted_response = f"""
                {response['answer']} Check out the Youtube video [{title}]({link}&t={int(start_time)}s)!
            """

        msg.content = formatted_response
    await msg.send()