import streamlit as st
from hazm import Normalizer, SentenceTokenizer
import os
import docx
from langchain.chat_models import ChatOpenAI
from langchain.schema import SystemMessage, HumanMessage
from rapidfuzz import fuzz

# ---------- استایل ----------
st.markdown("""
    <style>
    @import url('https://fonts.googleapis.com/css2?family=Vazirmatn:wght@400;700&display=swap');

    html, body, [class*="css"] {
        font-family: 'Vazirmatn', Tahoma, sans-serif;
        direction: rtl;
        text-align: right;
    }

    .stApp {
        background: linear-gradient(to left, #4b5e40, #2e3b2e);
        color: #ffffff;
    }

    [data-testid="stSidebar"] {
        width: 260px !important;
        background-color: #1a2b1e;
        border: none !important;
        padding-top: 20px;
    }

    .menu-item {
        display: flex;
        align-items: center;
        gap: 12px;
        padding: 12px 20px;
        font-size: 16px;
        color: #d4d4d4;
        cursor: pointer;
        transition: background-color 0.3s ease;
    }

    .menu-item:hover {
        background-color: #2e3b2e;
        color: #b8860b;
    }

    .menu-item img {
        width: 24px;
        height: 24px;
    }

    .stButton>button {
        background-color: #b8860b !important;
        color: #1a2b1e !important;
        font-family: 'Vazirmatn', Tahoma;
        font-weight: 700;
        border-radius: 10px;
        padding: 12px 24px;
        border: none;
        transition: all 0.3s ease;
        font-size: 16px;
        width: 100%;
        margin: 10px 0;
    }

    .stButton>button:hover {
        background-color: #8b6508 !important;
        transform: translateY(-2px);
        box-shadow: 0 4px 8px rgba(0,0,0,0.3);
    }

    .header-text {
        text-align: center;
        margin: 20px 0;
        background-color: rgba(26, 43, 30, 0.9);
        padding: 25px;
        border-radius: 15px;
        box-shadow: 0 6px 12px rgba(0,0,0,0.4);
    }

    .header-text h1 {
        font-size: 42px;
        color: #b8860b;
        margin: 0;
        font-weight: 700;
    }

    .subtitle {
        font-size: 18px;
        color: #d4d4d4;
        margin-top: 10px;
    }

    .chat-message {
        background-color: rgba(26, 43, 30, 0.95);
        border: 2px solid #b8860b;
        border-radius: 15px;
        padding: 20px;
        margin: 15px 0;
        box-shadow: 0 6px 12px rgba(0,0,0,0.3);
        animation: fadeIn 0.6s ease;
        font-size: 18px;
        color: #d4d4d4;
        display: flex;
        align-items: center;
        gap: 15px;
    }

    @keyframes fadeIn {
        from { opacity: 0; transform: translateY(10px); }
        to { opacity: 1; transform: translateY(0); }
    }

    .stTextInput>div>input, .stTextArea textarea {
        background-color: rgba(26, 43, 30, 0.95) !important;
        border-radius: 10px !important;
        border: 1px solid #b8860b !important;
        padding: 12px !important;
        font-family: 'Vazirmatn', Tahoma;
        font-size: 16px;
        color: #d4d4d4 !important;
    }

    hr {
        border: 1px solid #b8860b;
        margin: 15px 0;
    }

    [data-testid="stSidebar"] > div {
        border: none !important;
    }
    </style>
""", unsafe_allow_html=True)

# ---------- احراز هویت ----------
if "authenticated" not in st.session_state:
    st.session_state.authenticated = False

if not st.session_state.authenticated:
    st.markdown("<h3 style='text-align: center; color: #b8860b;'>ورود به رزم‌یار ارتش</h3>", unsafe_allow_html=True)
    username = st.text_input("نام کاربری:", placeholder="شناسه نظامی خود را وارد کنید")
    password = st.text_input("رمز عبور:", type="password", placeholder="رمز عبور نظامی")
    if st.button("ورود"):
        if username == "admin" and password == "123":
            st.session_state.authenticated = True
            st.rerun()
        else:
            st.error("نام کاربری یا رمز عبور اشتباه است.")
    st.stop()

# ---------- سایدبار ----------
with st.sidebar:
    st.image("log.png", use_container_width=True)
    menu_items = [
        ("گزارش عملیاتی", "https://cdn-icons-png.flaticon.com/512/3596/3596165.png"),
        ("تاریخچه ماموریت‌ها", "https://cdn-icons-png.flaticon.com/512/709/709496.png"),
        ("تحلیل داده‌های نظامی", "https://cdn-icons-png.flaticon.com/512/1828/1828932.png"),
        ("مدیریت منابع", "https://cdn-icons-png.flaticon.com/512/681/681494.png"),
        ("دستیار فرماندهی", "https://cdn-icons-png.flaticon.com/512/3601/3601646.png"),
        ("تنظیمات امنیتی", "https://cdn-icons-png.flaticon.com/512/2099/2099058.png"),
        ("پشتیبانی فنی", "https://cdn-icons-png.flaticon.com/512/597/597177.png"),
    ]
    for idx, (text, icon) in enumerate(menu_items):
        st.markdown(f"""
            <div class="menu-item">
                <img src="{icon}" />
                {text}
            </div>
        """, unsafe_allow_html=True)
        if idx in [1, 3, 5]:
            st.markdown("<hr/>", unsafe_allow_html=True)

# ---------- سربرگ ----------
st.markdown("""
    <div class="header-text">
        <h1>رزم‌یار ارتش</h1>
        <div class="subtitle">دستیار هوشمند ارتش جمهوری اسلامی ایران</div>
    </div>
""", unsafe_allow_html=True)

# ---------- مدل زبانی ----------
llm = ChatOpenAI(
    base_url="https://api.together.xyz/v1",
    api_key='0291f33aee03412a47fa5d8e562e515182dcc5d9aac5a7fb5eefdd1759005979',
    model="meta-llama/Llama-3.3-70B-Instruct-Turbo-Free"
)

# ---------- پردازش فایل‌ها ----------
folder_path = 'C:/Users/ici/Downloads/Telegram Desktop/45/46'
texts = []

for filename in os.listdir(folder_path):
    if filename.endswith(".docx"):
        full_path = os.path.join(folder_path, filename)
        doc = docx.Document(full_path)
        file_text = "\n".join([para.text for para in doc.paragraphs])
        if file_text.strip():
            texts.append(file_text)

normalizer = Normalizer()
sentence_tokenizer = SentenceTokenizer()

all_sentences = []
for text in texts:
    normalized = normalizer.normalize(text)
    sentences = sentence_tokenizer.tokenize(normalized)
    all_sentences.extend(sentences)

# ---------- ورودی جستجو ----------
query = st.text_input("🔎 کلمه یا عبارت موردنظر خود را وارد کنید:")

if query:
    found = False
    threshold = 80

    for idx, sentence in enumerate(all_sentences):
        similarity = fuzz.partial_ratio(query, sentence)
        if similarity >= threshold:
            next_sentences = []
            for i in range(1, 6):
                if idx + i < len(all_sentences):
                    next_sentences.append(all_sentences[idx + i])

            total_text = sentence + " " + " ".join(next_sentences)
            prompt = f"پاسخی که باید بازنویسی شود:\n{total_text}\n\nلطفاً این پاسخ را با در نظر گرفتن محتوای سوال زیر و لحن آن بازنویسی کن:\n\nسوال: {query}"
            response = llm([
                SystemMessage(content="You are a helpful assistant."),
                HumanMessage(content=prompt)
            ])
            rewritten = response.content.strip()
            st.markdown(f'<div class="chat-message">{rewritten}</div>', unsafe_allow_html=True)
            found = True
            break

    if not found:
        prompt = f"لطفاً بر اساس سوال زیر یک متن مرتبط و معنادار تولید کن:\n\nسوال: {query}"
        response = llm([
            SystemMessage(content="You are a helpful assistant."),
            HumanMessage(content=prompt)
        ])
        rewritten = response.content.strip()
        st.markdown(f'<div class="chat-message">{rewritten}</div>', unsafe_allow_html=True)