import streamlit as st
import requests
import base64
import os
import re

st.title("AI สนับสนุนความรู้ด้าน PDPA")
st.write("AI อัจฉริยะที่ช่วยการเรียนรู้ พระราชบัญญัติคุ้มครองข้อมูลส่วนบุคคล (PDPA) [Powered by NT Lahnmah]")

st.subheader("🏷️ Tags")
st.markdown("""
    - LLM   RAG   eRAG   PDPA   Machine Learning   AI   Data Privacy   GenAI   NT GenAI   ntgenai   lahnmah   NT Thai GPT   ntthaigpt   medical   medtech   HealthGPT   หลานม่า   NT Academy
""", unsafe_allow_html=True)

system_prompt = "คุณเป็นผู้ช่วยที่มีความรู้ด้านกฎหมาย พระราชบัญญัติคุ้มครองข้อมูลส่วนบุคคล (PDPA) และสามารถให้คำตอบที่เกี่ยวข้องเฉพาะตาม context ที่ได้รับ"

def clean_text_for_search(text):
    return ' '.join(re.sub(r'P-\d+\s*$', '', text, flags=re.MULTILINE).split())

def create_highlighted_pdf(pdf_path, search_text, page_number):
    try:
        search_text = clean_text_for_search(search_text)
        doc = fitz.open(pdf_path)
        page = doc[int(page_number) - 1]
        words = [word for word in search_text.split() if word]

        for word in words:
            if len(word) > 3:
                text_instances = page.search_for(word)
                for inst in text_instances:
                    highlight = page.add_highlight_annot(inst)
                    highlight.set_colors(stroke=(1, 1, 0))
                    highlight.update()

        new_doc = fitz.open()
        new_doc.insert_pdf(doc, from_page=int(page_number) - 1, to_page=int(page_number) - 1)
        pdf_bytes = new_doc.write()

        doc.close()
        new_doc.close()

        return pdf_bytes
    except Exception as e:
        st.error(f"Error in create_highlighted_pdf: {str(e)}")
        return None

def format_file_size(size_in_bytes):
    for unit in ['B', 'KB', 'MB', 'GB']:
        if size_in_bytes < 1024:
            return f"{size_in_bytes:.2f} {unit}"
        size_in_bytes /= 1024
    return f"{size_in_bytes:.2f} GB"

def display_search_result(result, index):
    with st.expander(f"🔍 Search Result #{index + 1} (Score: {result['score']:.4f})"):
        st.markdown("#### 📄 Document Information")
        col1, col2 = st.columns(2)

        with col1:
            st.markdown("**File Details:**")
            st.write(f"• File Name: {result['metadata']['file_name']}")
            st.write(f"• Page: {result['metadata']['page_label']}")
            st.write(f"• Type: {result['metadata']['file_type']}")
            st.write(f"• Size: {format_file_size(result['metadata']['file_size'])}")

        with col2:
            st.markdown("**Dates:**")
            st.write(f"• Created: {result['metadata']['creation_date']}")
            st.write(f"• Modified: {result['metadata']['last_modified_date']}")

        st.markdown("#### 📝 Content")
        st.markdown(f"```\n{result['text']}\n```")

        try:
            pdf_path = result['file_path']
            if os.path.exists(pdf_path):
                st.markdown("#### 📄 PDF Preview (with highlighted text)")

                highlighted_pdf = create_highlighted_pdf(
                    pdf_path,
                    result['text'],
                    result['metadata']['page_label']
                )

                if highlighted_pdf:
                    base64_pdf = base64.b64encode(highlighted_pdf).decode('utf-8')
                    pdf_display = f'''
                        <iframe 
                            src="data:application/pdf;base64,{base64_pdf}" 
                            width="100%" 
                            height="800px" 
                            type="application/pdf"
                            style="border: 1px solid #ccc; border-radius: 5px;"
                        ></iframe>
                    '''
                    st.markdown(pdf_display, unsafe_allow_html=True)
                else:
                    st.error("Failed to create highlighted PDF")
        except Exception as e:
            st.error(f"Error displaying PDF: {str(e)}")


with st.form(key="input_form"):
    user_input = st.text_input("You:", value="ข้อมูลส่วนบุคคล คืออะไร มีกี่ประเภท", key="input")
    submit_button = st.form_submit_button("Send")

if submit_button and user_input:
    try:
        response = requests.post("http://113.53.253.50:8002/search", json={"query": user_input})
        response.raise_for_status()
        data = response.json()
        search_results = data["results"]

        st.markdown("### 🔎 Search Results")
        for idx, result in enumerate(search_results):
            display_search_result(result, idx)

    except requests.RequestException as e:
        st.error(f"Error: {str(e)}")


st.subheader("📄 ไฟล์เอกสารที่ใช้เป็นคลังข้อมูล PDPA")
st.markdown(
    """
    <a href="https://huggingface.co/spaces/amornpan/eRAG-PDPA-v1/tree/main/%E0%B9%80%E0%B8%AD%E0%B8%81%E0%B8%AA%E0%B8%B2%E0%B8%A3_PDPA" style="text-decoration: none;">
        <span style="font-size: 24px;">👉 ค้นหาเฉพาะไฟล์เหล่านี้เท่านั้น!</span>
    </a>
    """,
    unsafe_allow_html=True
)

st.subheader("👤 Authors")
st.write("""
- Amornpan Phornchaicharoen (amornpan@gmail.com)
- Aekanun Thongtae (cto@bangkokfirsttech.com)
- Montita Somsoo (montita.fonn@gmail.com)
- Jiranuwat Songpad (jiranuwat.song64@gmail.com)
- Phongsatorn Somjai (ipongdev@gmail.com)
- Benchawan Wangphoomyai (benchaa.27@gmail.com)
""")