import time
import json
import requests
import streamlit as st


st.set_page_config(page_title="ViBidLQA - Trợ lý AI hỗ trợ hỏi đáp luật Việt Nam", page_icon="./app/static/ai.jpg", layout="wide", initial_sidebar_state="expanded")

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

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. 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/ai.jpg"/>
</div>
""", unsafe_allow_html=True)
st.markdown("<h2 style='text-align: center;'>The ViBidLQA System </h2>", unsafe_allow_html=True)

url_api_retrieval_model = st.sidebar.text_input(label="URL API Retrieval model:")
url_api_extraction_model = st.sidebar.text_input(label="URL API Extraction model:")
url_api_generation_model = st.sidebar.text_input(label="URL API Generation model:")

answering_method = st.sidebar.selectbox(options=['Extraction', 'Generation'], label='Chọn mô hình trả lời câu hỏi:', index=0)

if answering_method == 'Generation':
    print('Switching to generative model...')
    print('Loading generative model...')

if answering_method == 'Extraction':
    print('Switching to extraction model...')
    print('Loading extraction model...')

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"]
        print(f"Văn bản pháp luật được truy hồi: {results[0]['text']}")
        print("="*100)
        return results[0]["text"]
    else:
        return f"Lỗi: {response.status_code} - {response.text}"

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

    response = requests.post(url_api_generation_model, json=data)
    if response.status_code == 200:
        result = response.json()
        return result["answer"]
    else:
        return f"Lỗi: {response.status_code} - {response.text}"

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

    # Sử dụng requests với stream=True
    response = requests.post(url_api_generation_model, json=data, stream=True)
    
    if response.status_code == 200:
        # Trả về response để xử lý streaming
        return response
    else:
        return f"Lỗi: {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])

def get_extractive_answer(question, stride=20, max_length=256, n_best=50, max_answer_length=512):
    context = retrieve_context(question=question)
    
    data = {
        "context": context,
        "question": question,
        "stride": stride,
        "max_length": max_length,
        "n_best": n_best,
        "max_answer_length": max_answer_length
    }

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

    if response.status_code == 200:
        result = response.json()
        return result["best_answer"]
    else:
        return f"Lỗi: {response.status_code} - {response.text}"

for message in st.session_state.messages:
    if message['role'] == 'assistant':
        avatar_class = "assistant-avatar"
        message_class = "assistant-message"
        avatar = './app/static/ai.jpg'
    else:
        avatar_class = "user-avatar"
        message_class = "user-message"
        avatar = './app/static/human.png'
    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">
        <img src="./app/static/human.png" class="user-avatar" />
        <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 = ""
    if answering_method == 'Generation':
        response_stream = get_abstractive_answer_stream(question=prompt)
        
        if isinstance(response_stream, str):
            full_response = response_stream
            message_placeholder.markdown(f"""
            <div class="assistant-message">
                <img src="./app/static/ai.jpg" class="assistant-avatar" />
                <div class="stMarkdown">{full_response}</div>
            </div>
            """, unsafe_allow_html=True)
        else:
            full_response = ""
            for line in response_stream.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/ai.jpg" class="assistant-avatar" />
                                <div class="stMarkdown">{full_response}●</div>
                            </div>
                            """, unsafe_allow_html=True)
                            
                        except json.JSONDecodeError:
                            pass
    else:
        ext_answer = get_extractive_answer(question=prompt)
        for word in generate_text_effect(ext_answer):
            full_response = word
            
            message_placeholder.markdown(f"""
            <div class="assistant-message">
                <img src="./app/static/ai.jpg" class="assistant-avatar" />
                <div class="stMarkdown">{full_response}●</div>
            </div>
            """, unsafe_allow_html=True)

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