from openai import OpenAI
import crawl_the_site
import create_database
import streamlit as st
import time

api_key = "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"

openai_client = OpenAI(base_url="https://openrouter.ai/api/v1",api_key=api_key)

prompt_template = """Answer the question based only on the information provided below.Instead of directing the user to resources, answer the user using the content of the resources.And dont say anything like "Based on the provided information" etc. in the answer.Provide examples based on the users question.For example, when he says "how do i create a slash command" , provide him with an example.
## Information : 
{}

## Question : 
# {}"""

if not "messages" in st.session_state:
    st.session_state.messages = []
if not "stage" in st.session_state:
    st.session_state.stage = "crawl-input"
elif st.session_state.stage == "remove-crawl-input-widgets":
    st.session_state.stage = "prompt-input"

if st.session_state.stage == "crawl-input":
    config = st.text_area("Enter the start URL and the limit of crawling (in this format : URL,limit) :")
    config_btn = st.button("Start crawling!")

    if config and config_btn:
        start = config.split(",")[0]
        limit = int(config.split(",")[1])
            
        txt = crawl_the_site.crawl(start,limit)
        st.session_state.collection = create_database.create_database(txt)
        st.write("Crawling is done. Click \"Start Crawling!\" again to start chatting!")
        st.session_state.stage = "remove-crawl-input-widgets"

if st.session_state.stage == "prompt-input":
    
    prompt = st.chat_input("Ask the documentation a question!")
                
    if prompt:
        for message in st.session_state.messages:
            with st.chat_message(message["role"]):
                st.markdown(message["content"])
        
        with st.chat_message("user"):
            st.markdown(prompt)
                
        st.session_state.messages.append({"role":"user","content":prompt})
        results = st.session_state.collection.query(query_texts=[prompt],n_results=10)
        infos = results["documents"][0]
        
        info_text = ""
        for info in infos:
            info_text += info + "\n-----\n"
        info_text = info_text.strip()
        
        prompt = prompt_template.format(info_text,prompt)
        
        completion = openai_client.chat.completions.create(
        extra_headers={},
        extra_body={},
        model="deepseek/deepseek-chat-v3-0324:free",
        messages=[{"role":"user","content":prompt}])
        
        with st.chat_message("assistant"):
            full_response = ""
            message_placeholder = st.empty()
            for chunk in completion.choices[0].message.content.split():
                full_response += chunk + " "
                time.sleep(0.05)
                message_placeholder.markdown(full_response + "▌")
            message_placeholder.markdown(full_response)
        st.session_state.messages.append({"role":"assistant","content":full_response})