import os

import chainlit as cl
import dotenv
from chainlit import CustomElement
from llama_index.core import StorageContext, load_index_from_storage
from vdb import further_retrieve, split_markdown_by_section, build_documents, create_vector_index, completion
import re

# storage_context = StorageContext.from_defaults(persist_dir=f'./md_vector_store')
# index = load_index_from_storage(storage_context)

md_file = "opertgfr2.md"

sections = split_markdown_by_section(md_file)
docs = build_documents(sections)
index = create_vector_index(docs)

dotenv.load_dotenv()


@cl.set_starters
async def set_starters():
    questions = [
        "C'est quoi les courses et origines",
        "Accès aux attributs de personnelisation du système",
        "Mise sous tension/arrêt",
        # "Qui signe les bon de commande pour superieur de 80000 euros pour ordinateur",
        # "Qui signe les bon de commande pour tablette pour les ecoles",
        # "Qui signe marchés à procédure adaptée pour 90000 euros pour des plantes",
        # "Qui signe marchés à procédure adaptée pour 40000 euros",
        # "qui signe les bdc a 25k pour Achat de matériel comptable",
        # "Qui signe un bdc d'achat de lampadaire de rue pour 20000 euros ?",
    ]

    return [
        cl.Starter(
            label=value,
            message=value
        ) for value in questions
    ]


@cl.on_chat_start
async def on_start():

    return

"""https://aqjnuiodpkwdplgeozck.supabase.co/storage/v1/object/public/faq_chunker/images/042eb615b98ba399a470268106a7319da3b4dad9e518bef62182d0e4fdc30988.png"""
@cl.on_message
async def main(message):
    is_first = cl.user_session.get("is_first", True)
    if is_first:
        await cl.Message("", elements=[CustomElement(name="ImageGallery")]).send()
        cl.user_session.set("is_first", False)
    messages = cl.user_session.get("messages", [])
    response = await further_retrieve(message.content, index, messages)
    msg = cl.Message("")
    message_response = ""
    async for token in response:
        message_response += token

        message_response = re.sub(r"\!\[.*?\]\(images\/(.*)?\)",
                                # r"""<img classname="m-2 w-full" src="https://aqjnuiodpkwdplgeozck.supabase.co/storage/v1/object/public/faq_chunker/images/\1">""",
                                  r"![](https://aqjnuiodpkwdplgeozck.supabase.co/storage/v1/object/public/faq_chunker/images/\1)",
                                  message_response)
        await msg.stream_token(message_response, is_sequence=True)
    await msg.update()
    messages.extend(
        [
            {"role": "user", "content": message.content},
            {"role": "assistant", "content": message_response},
        ]
    )

    cl.user_session.set("messages", messages)