import time
import json
import requests
import streamlit as st
import os
from urllib.parse import urlencode, urlparse, parse_qs

st.set_page_config(page_title="IntelliChat - Trợ lý AI văn bản pháp luật Việt Nam", page_icon="./app/static/IntelliChat.png", layout="centered", initial_sidebar_state="collapsed")

# ==== MÔI TRƯỜNG OAuth ====
FB_APP_ID = os.getenv("FB_APP_ID")
FB_APP_SECRET = os.getenv("FB_APP_SECRET")
FB_REDIRECT_URI = os.getenv("FB_REDIRECT_URI")
FB_CLIENT_URL = os.getenv("FB_CLIENT_URL", "https://www.facebook.com")
FB_API_URL = os.getenv("FB_API_URL", "https://graph.facebook.com")
FB_BACKEND_URL = os.getenv("FB_BACKEND_URL")

# ==== MODULE URL ====
routing_response_module = st.secrets["ViBidLQA_Routing_Module"]
retrieval_module = st.secrets["ViBidLQA_Retrieval_Module"]
reranker_module = st.secrets["ViBidLQA_Rerank_Module"]
abs_QA_module = st.secrets["ViBidLQA_AQA_Module"]


url_api_question_classify_model = f"{routing_response_module}/query_classify"
url_api_unrelated_question_response_model = f"{routing_response_module}/response_unrelated_question"
url_api_introduce_system_model = f"{routing_response_module}/about_me"
url_api_retrieval_model = f"{retrieval_module}/search"
url_api_reranker_model = f"{reranker_module}/rerank"
url_api_generation_model = f"{abs_QA_module}/answer"
url_api_extract_reference_model = f"{routing_response_module}/extract_references_unstream"

# ========== STREAMLIT UI ==========

with open("./static/styles.css") as f:
    st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)

# ==== GIAO DIỆN CHÍNH - TABS ====
# tab1, tab2 = st.tabs(["🤖 ViBidLQA Chatbot", "🔐 Facebook OAuth"])

# =============================
# TAB 1: VIBIDLQA CHATBOT
# =============================
# with tab1:
if 'messages' not in st.session_state:
    st.session_state.messages = [{'role': 'assistant', 'content': "Xin chào. Tôi là trợ lý AI văn bản luật Đấu thầu Việt Nam được phát triển bởi Nguyễn Trường Phúc - sinh viên năm cuối ngành Khoa học Máy tính tại trường Đại học Sư phạm Kỹ thuật Hưng Yên. Rất vui khi được hỗ trợ bạn trong các vấn đề pháp lý tại Việt Nam!"}]

# st.markdown(f"""
# <div class=logo_area>
#     <img src="./app/static/IntelliChat.png"/>
# </div>
# """, unsafe_allow_html=True)
# st.markdown("<h2 style='text-align: center;'>ViBidLQA</h2>", unsafe_allow_html=True)

def classify_question(question):
    data = {
        "question": question
    }
    
    response = requests.post(url_api_question_classify_model, json=data)
    
    if response.status_code == 200:
        print(response)
        return response
    else:
        return f"Lỗi: {response.status_code} - {response.text}"

def introduce_system(question):
    data = {
        "question": question
    }

    response = requests.post(url_api_introduce_system_model, json=data, stream=True)

    if response.status_code == 200:
        return response
    else:
        return f"Lỗi: {response.status_code} - {response.text}"
    
def response_unrelated_question(question):
    data = {
        "question": question
    }

    response = requests.post(url_api_unrelated_question_response_model, json=data, stream=True)

    if response.status_code == 200:
        return response
    else:
        return f"Lỗi: {response.status_code} - {response.text}"

def retrieve_context(question, top_k=10):
    data = {
        "query": question,
        "top_k": top_k
    }

    response = requests.post(url_api_retrieval_model, json=data)

    if response.status_code == 200:
        results = response.json()["results"]
        return results
    else:
        return f"Lỗi tại Retrieval Module: {response.status_code} - {response.text}"

def rerank_context(url_rerank_module, question, relevant_docs, top_k=5):
    data = {
        "question": question,
        "relevant_docs": relevant_docs,
        "top_k": top_k
    }

    response = requests.post(url_rerank_module, json=data)

    if response.status_code == 200:
        results = response.json()["reranked_docs"]
        return results
    else:
        return f"Lỗi tại Rerank module: {response.status_code} - {response.text}"

def get_abstractive_answer(context, question):
    data = {
        "context": context,
        "question": question
    }

    response = requests.post(url_api_generation_model, json=data, stream=True)

    if response.status_code == 200:
        return response
    else:
        return f"Lỗi: {response.status_code} - {response.text}"

def get_references(context, question, answer):
    data = {
        "context": context,
        "question": question,
        "answer": answer
    }

    response = requests.post(url_api_extract_reference_model, json=data)

    if response.status_code == 200:
        return response.json()["refs"]
    else:
        return f"Lỗi tại module Reference Extractor: {response.status_code} - {response.text}"

def generate_text_effect(answer):
    words = answer.split()
    for i in range(len(words)):
        time.sleep(0.03)
        yield " ".join(words[:i+1])

for message in st.session_state.messages:
    if message['role'] == 'assistant':
        avatar_class = "assistant-avatar"
        message_class = "assistant-message"
        avatar = './app/static/IntelliChat.png'
    else:
        avatar_class = ""
        message_class = "user-message"
        avatar = ''
    st.markdown(f"""
    <div class="{message_class}">
        <img src="{avatar}" class="{avatar_class}" />
        <div class="stMarkdown">{message['content']}</div>
    </div>
    """, unsafe_allow_html=True)

if prompt := st.chat_input(placeholder='Tôi có thể giúp được gì cho bạn?'):
    st.markdown(f"""
    <div class="user-message">
            <div class="stMarkdown">{prompt}</div>
    </div>
    """, unsafe_allow_html=True)
    st.session_state.messages.append({'role': 'user', 'content': prompt})
    
    message_placeholder = st.empty()
    
    full_response = ""
    classify_result = classify_question(question=prompt).json()

    print(f"The type of user query: {classify_result}")

    if classify_result == "BIDDING_RELATED":
        retrieved_context = retrieve_context(question=prompt, top_k=10)
        retrieved_context = [item['text'] for item in retrieved_context]
        reranked_context = rerank_context(url_rerank_module=url_api_reranker_model,
                                      question=prompt,
                                      relevant_docs=retrieved_context,
                                      top_k=5)[0]
        
        abs_answer = get_abstractive_answer(context=reranked_context, question=prompt)

        if isinstance(abs_answer, str):
            full_response = abs_answer
            message_placeholder.markdown(f"""
            <div class="assistant-message">
                <img src="./app/static/IntelliChat.png" class="assistant-avatar" />
                <div class="stMarkdown">{full_response}</div>
            </div>
            """, unsafe_allow_html=True)
        else:
            full_response = ""
            for line in abs_answer.iter_lines():
                if line:
                    line = line.decode('utf-8')
                    if line.startswith('data: '):
                        data_str = line[6:]
                        if data_str == '[DONE]':
                            break
                        
                        try:
                            data = json.loads(data_str)
                            token = data.get('token', '')
                            full_response += token
                            
                            message_placeholder.markdown(f"""
                            <div class="assistant-message">
                                <img src="./app/static/IntelliChat.png" class="assistant-avatar" />
                                <div class="stMarkdown">{full_response}●</div>
                            </div>
                            """, unsafe_allow_html=True)
                            
                        except json.JSONDecodeError:
                            pass
        
            refs = st.expander("Tài liệu tham khảo", expanded=False)
            refs_list = get_references(context=reranked_context, question=prompt, answer=full_response)
            print(refs_list)
            refs.write(f"{refs_list}")

    elif classify_result == "ABOUT_CHATBOT":
        answer = introduce_system(question=prompt)

        if isinstance(answer, str):
            full_response = answer
            message_placeholder.markdown(f"""
            <div class="assistant-message">
                <img src="./app/static/IntelliChat.png" class="assistant-avatar" />
                <div class="stMarkdown">{full_response}</div>
            </div>
            """, unsafe_allow_html=True)
        else:
            full_response = ""
            for line in answer.iter_lines():
                if line:
                    line = line.decode('utf-8')
                    if line.startswith('data: '):
                        data_str = line[6:]
                        if data_str == '[DONE]':
                            break
                        
                        try:
                            data = json.loads(data_str)
                            token = data.get('token', '')
                            full_response += token
                            
                            message_placeholder.markdown(f"""
                            <div class="assistant-message">
                                <img src="./app/static/IntelliChat.png" class="assistant-avatar" />
                                <div class="stMarkdown">{full_response}●</div>
                            </div>
                            """, unsafe_allow_html=True)
                            
                        except json.JSONDecodeError:
                            pass
    
    else:
        answer = response_unrelated_question(question=prompt)

        if isinstance(answer, str):
            full_response = answer
            message_placeholder.markdown(f"""
            <div class="assistant-message">
                <img src="./app/static/IntelliChat.png" class="assistant-avatar" />
                <div class="stMarkdown">{full_response}</div>
            </div>
            """, unsafe_allow_html=True)
        else:
            full_response = ""
            for line in answer.iter_lines():
                if line:
                    line = line.decode('utf-8')
                    if line.startswith('data: '):
                        data_str = line[6:]
                        if data_str == '[DONE]':
                            break
                        
                        try:
                            data = json.loads(data_str)
                            token = data.get('token', '')
                            full_response += token
                            
                            message_placeholder.markdown(f"""
                            <div class="assistant-message">
                                <img src="./app/static/IntelliChat.png" class="assistant-avatar" />
                                <div class="stMarkdown">{full_response}●</div>
                            </div>
                            """, unsafe_allow_html=True)
                            
                        except json.JSONDecodeError:
                            pass   

    message_placeholder.markdown(f"""
    <div class="assistant-message">
        <img src="./app/static/IntelliChat.png" class="assistant-avatar" />
            <div class="stMarkdown">
                {full_response}
            </div>
    </div>
    """, unsafe_allow_html=True)
    
    st.session_state.messages.append({'role': 'assistant', 'content': full_response})

# =============================
# TAB 2: FACEBOOK OAUTH
# =============================
# with tab2:
#     st.title("Đăng nhập Facebook để lấy Page Access Token")

#     # Tạo link login
#     login_url = f"{FB_BACKEND_URL}/facebook/login"
    
#     st.markdown(f"[👉 Bấm vào đây để đăng nhập Facebook]({login_url})")
    
#     st.info("Sau khi đăng nhập xong, bạn có thể quay lại ứng dụng này. Thông tin page đã được in ra ở backend.")