import gradio as gr
import re
import os
import requests
import time
import logging

# تنظیم logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class PreciseAnonymizer:
    def __init__(self):
        self.mapping_table = {}
        self.counters = {
            'company': 0, 
            'person': 0, 
            'amount': 0, 
            'percent': 0
        }
        self.api_key = os.getenv("OPENAI_API_KEY", "")

    def anonymize_text(self, original_text, lang='fa'):
        """ناشناس‌سازی دقیق براساس مثال‌های واقعی"""
        try:
            if not original_text or not original_text.strip():
                return "❌ Please enter input text!" if lang == 'en' else "❌ لطفاً متن ورودی را وارد کنید!"
            
            # ریست متغیرها
            self.mapping_table = {}
            self.counters = {key: 0 for key in self.counters.keys()}
            
            anonymized = original_text
            
            # الگوهای دقیق براساس مثال‌های واقعی
            patterns = [
                # شرکت‌ها - الگوهای مشخص
                (r'شرکت سرمایه‌گذاری دارویی تأمین \(تیپیکو\)', 'company'),
                (r'شرکت پتروشیمی بوعلی سینا', 'company'),
                (r'شرکت پتروشیمی پارس', 'company'),
                (r'شرکت آسان پادرو', 'company'),
                (r'شرکت پالایش نفت اصفهان', 'company'),
                (r'شرکت گروه توسعه مالی مهر آیندگان \(ومهان\)', 'company'),
                (r'شرکت وانیا نیک تدبیر', 'company'),
                (r'تدوین و همکاران', 'company'),
                (r'سازمان حسابرسی', 'company'),
                (r'سرزمین هوشمند پاد \(زیرمجموعه بانک پاسارگاد\)', 'company'),
                (r'ایران خودرو', 'company'),
                (r'بانک پاسارگاد', 'company'),
                (r'بانک ملت', 'company'),
                (r'بانک سرمایه', 'company'),
                (r'همراه اول', 'company'),
                (r'گروه همراه اول', 'company'),
                
                # نام اشخاص با عناوین
                (r'مهدی اخوان بهابادی، مدیرعامل همراه اول', 'person'),
                (r'فرج‌اله قدمی', 'person'),
                
                # مبالغ مالی - الگوهای دقیق
                (r'\d+,\d{3} میلیارد ریال', 'amount'),  # 681,667 میلیارد ریال
                (r'\d+ هزار و \d+ میلیارد تومان', 'amount'),  # 23 هزار و 296 میلیارد تومان
                (r'\d+ هزار و \d+ دستگاه', 'amount'),  # 537 هزار و 736 دستگاه
                (r'بیش از \d+ همت', 'amount'),  # بیش از 37 همت
                (r'حدود \d+ میلیون تومان', 'amount'),  # حدود 69 میلیون تومان
                (r'بیش از \d+ همت', 'amount'),  # بیش از 111 همت
                (r'\d+,\d{3},\d{3} میلیارد ریال', 'amount'),  # 681,667 میلیارد ریال
                (r'\d+,\d{3} میلیارد ریال', 'amount'),  # 86,278 میلیارد ریال
                (r'\d+,\d{3} میلیارد ریال', 'amount'),  # 12,140 میلیارد ریال
                (r'\d+,\d{3} میلیارد ریال', 'amount'),  # 51,670 میلیارد ریال
                (r'\d+,\d{3} میلیارد ریال', 'amount'),  # 33,000 میلیارد ریال
                (r'\d+ هزار میلیارد ریالی', 'amount'),  # 155 هزار میلیارد ریالی
                (r'\d+ میلیارد تومانی', 'amount'),  # 2700 میلیارد تومانی
                (r'نزدیک به \d+ هزار میلیارد تومان', 'amount'),  # نزدیک به 67 هزار میلیارد تومان
                (r'بیش از \d+ میلیارد تومان', 'amount'),  # بیش از 6 میلیارد تومان
                (r'حدود \d+ میلیارد تومان', 'amount'),  # حدود 30 میلیارد تومان
                (r'حدود \d+ میلیون دلار', 'amount'),  # حدود 100 میلیون دلار
                (r'\d+ تا \d+\.\d+ میلیون تن', 'amount'),  # 1 تا 1.5 میلیون تن
                (r'\d+ ریال', 'amount'),  # 936 ریال
                (r'\d+ ریال', 'amount'),  # 446 ریال
                (r'\d+ ریال', 'amount'),  # 610 ریال
                (r'\d+ هزار و \d+ میلیارد تومان', 'amount'),  # 70 هزار و 677 میلیارد تومان
                (r'\d+ میلیارد تومان', 'amount'),  # 7101 میلیارد تومان
                (r'\d+ میلیارد تومان', 'amount'),  # 8003 میلیارد تومان
                (r'\d+ هزار میلیارد تومان', 'amount'),  # 49 هزار میلیارد تومان
                (r'\d+ میلیارد تومانی', 'amount'),  # 178 میلیارد تومانی
                (r'\d+\.\d+ میلیون نفر', 'amount'),  # 73.7 میلیون نفر
                (r'\d+ هزار و \d+ میلیارد تومان', 'amount'),  # 16 هزار و 495 میلیارد تومان
                (r'\d+ هزار و \d+ میلیارد تومان', 'amount'),  # 31 هزار و 756 میلیارد تومان
                
                # درصدها - الگوهای دقیق
                (r'\d+\.\d+ درصد', 'percent'),  # 4.58 درصد
                (r'\d+ درصد', 'percent'),  # 75 درصد
                (r'منفی \d+ درصد', 'percent'),  # منفی 345 درصد
                (r'\d+ درصد', 'percent'),  # 97 درصد
                (r'بیش از \d+ درصد', 'percent'),  # بیش از 60 درصد
                (r'\d+ درصد', 'percent'),  # 37 درصد
                (r'\d+ درصد', 'percent'),  # 15 درصد
                (r'\d+ درصد', 'percent'),  # 14 درصد
                (r'\d+ الی \d+ درصد', 'percent'),  # 50 الی 70 درصد
                (r'\d+ درصدی', 'percent'),  # 14 درصدی
                (r'\d+ درصدی', 'percent'),  # 37 درصدی
                (r'\d+ درصدی', 'percent'),  # 15 درصدی
            ]
            
            # پردازش الگوها به ترتیب از طولانی‌ترین به کوتاه‌ترین
            for pattern, category in patterns:
                matches = list(re.finditer(pattern, anonymized, re.IGNORECASE))
                
                # مرتب‌سازی matches بر اساس طول (طولانی‌ترین اول)
                matches.sort(key=lambda x: len(x.group(0)), reverse=True)
                
                for match in matches:
                    matched_text = match.group(0)
                    
                    # بررسی که قبلاً جایگزین نشده باشد
                    if matched_text in anonymized and matched_text not in self.mapping_table:
                        self.counters[category] += 1
                        code = f"{category}-{self.counters[category]}"
                        self.mapping_table[matched_text] = code
                        anonymized = anonymized.replace(matched_text, code)
                        logger.info(f"Replaced: {matched_text} -> {code}")
            
            logger.info(f"✅ Anonymization completed. Found {len(self.mapping_table)} entities.")
            return anonymized
            
        except Exception as e:
            return f"❌ Error in anonymization: {str(e)}" if lang == 'en' else f"❌ خطا در ناشناس‌سازی: {str(e)}"

    def send_to_chatgpt(self, anonymized_text, lang='fa'):
        """ارسال به ChatGPT"""
        try:
            if not anonymized_text or not anonymized_text.strip():
                return "❌ Anonymized text is empty!" if lang == 'en' else "❌ متن ناشناس‌شده خالی است!"
            
            if not self.api_key:
                return "❌ API Key not configured! Please set OPENAI_API_KEY environment variable." if lang == 'en' else "❌ کلید API تنظیم نشده است! لطفاً OPENAI_API_KEY را در متغیرهای محیطی تنظیم کنید."
            
            system_msg = "You are a professional financial analyst. The text contains anonymous codes. Answer questions accurately." if lang == 'en' else "شما یک تحلیلگر مالی حرفه‌ای هستید. متن حاوی کدهای ناشناس است. به سوالات با دقت پاسخ دهید."
            
            headers = {
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json"
            }
            
            data = {
                "model": "gpt-4o-mini",
                "messages": [
                    {"role": "system", "content": system_msg},
                    {"role": "user", "content": anonymized_text}
                ],
                "max_tokens": 2000,
                "temperature": 0.7
            }
            
            response = requests.post(
                "https://api.openai.com/v1/chat/completions",
                headers=headers,
                json=data,
                timeout=30
            )
            
            if response.status_code == 200:
                result = response.json()
                return result['choices'][0]['message']['content']
            else:
                error_data = response.json() if response.content else {}
                error_message = error_data.get('error', {}).get('message', response.text)
                
                if 'Incorrect API key' in error_message:
                    return "❌ Invalid API key." if lang == 'en' else "❌ کلید API نامعتبر است."
                elif 'quota' in error_message:
                    return "❌ API quota exceeded." if lang == 'en' else "❌ سهمیه API تمام شده است."
                else:
                    return f"❌ API Error: {error_message}"
                
        except Exception as e:
            return f"❌ Error connecting to ChatGPT: {str(e)}" if lang == 'en' else f"❌ خطا در ارتباط با ChatGPT: {str(e)}"

    def deanonymize_response(self, gpt_response, lang='fa'):
        """بازگردانی"""
        try:
            if not gpt_response or not gpt_response.strip():
                return "❌ ChatGPT response is empty!" if lang == 'en' else "❌ پاسخ ChatGPT خالی است!"
            
            if not self.mapping_table:
                return "❌ Mapping table is empty!" if lang == 'en' else "❌ جدول نگاشت خالی است!"
            
            final_result = gpt_response
            reverse_mapping = {code: original for original, code in self.mapping_table.items()}
            
            # جایگزینی از طولانی‌ترین کد اول
            sorted_codes = sorted(reverse_mapping.items(), key=lambda x: len(x[0]), reverse=True)
            for code, original in sorted_codes:
                final_result = final_result.replace(code, original)
            
            return final_result
            
        except Exception as e:
            return f"❌ Deanonymization error: {str(e)}" if lang == 'en' else f"❌ خطا در بازگردانی: {str(e)}"

def process_all_steps(input_text, language):
    """پردازش خودکار تمام مراحل"""
    lang = 'en' if language == 'English' else 'fa'
    
    if not input_text.strip():
        error_msg = "❌ Please enter input text!" if lang == 'en' else "❌ لطفاً متن ورودی را وارد کنید!"
        return error_msg, "", "", ""
    
    try:
        start_time = time.time()
        
        anonymized_text = anonymizer.anonymize_text(input_text, lang)
        if anonymized_text.startswith("❌"):
            return anonymized_text, "", "", ""
        
        gpt_response = anonymizer.send_to_chatgpt(anonymized_text, lang)
        if gpt_response.startswith("❌"):
            entities_found = len(anonymizer.mapping_table)
            success_msg = (f"✅ Anonymization completed!\n"
                          f"📊 Total: {entities_found} entities protected")
            return success_msg, anonymized_text, gpt_response, ""
        
        final_result = anonymizer.deanonymize_response(gpt_response, lang)
        
        total_time = time.time() - start_time
        entities_found = len(anonymizer.mapping_table)
        
        # آمار تفصیلی
        company_count = anonymizer.counters['company']
        amount_count = anonymizer.counters['amount']
        percent_count = anonymizer.counters['percent']
        person_count = anonymizer.counters['person']
        
        success_msg = (f"🎉 Complete anonymization & restoration successful!\n"
                      f"🏢 Companies: {company_count} | 💰 Amounts: {amount_count} | 📊 Percentages: {percent_count} | 👤 Persons: {person_count}\n"
                      f"📊 Total: {entities_found} entities | ⏱️ Time: {total_time:.2f}s")
        
        return success_msg, anonymized_text, gpt_response, final_result
        
    except Exception as e:
        error_msg = f"❌ Processing error: {str(e)}" if lang == 'en' else f"❌ خطا در پردازش: {str(e)}"
        return error_msg, "", "", ""

def get_mapping_table(language):
    """نمایش جدول نگاشت"""
    lang = 'en' if language == 'English' else 'fa'
    
    if not anonymizer.mapping_table:
        return "❌ Mapping table is empty! Please process some text first." if lang == 'en' else "❌ جدول نگاشت خالی است! ابتدا متنی را پردازش کنید."
    
    result = "📋 **Precise Mapping Table:**\n\n" if lang == 'en' else "📋 **جدول نگاشت دقیق:**\n\n"
    
    # گروه‌بندی بر اساس نوع
    categories = {
        'company': '🏢 **Companies**',
        'amount': '💰 **Amounts**',
        'percent': '📊 **Percentages**',
        'person': '👤 **Persons**'
    }
    
    for category, title in categories.items():
        category_items = {k: v for k, v in anonymizer.mapping_table.items() if v.startswith(category)}
        if category_items:
            result += f"{title}:\n"
            for original, code in category_items.items():
                result += f"   • `{original}` → `{code}`\n"
            result += "\n"
    
    # آمار کلی
    result += f"📊 **Summary**: {len(anonymizer.mapping_table)} total entities anonymized\n"
    
    return result

def clear_all():
    """پاک کردن همه"""
    anonymizer.mapping_table = {}
    anonymizer.counters = {key: 0 for key in anonymizer.counters.keys()}
    return "", "", "", "", ""

def update_ui_text(language):
    """به‌روزرسانی متن‌های رابط کاربری"""
    if language == 'English':
        return {
            'title': 'Precise Business Data Anonymization System',
            'step1': 'Input Text & Settings',
            'step2': 'Anonymized Text',
            'step3': 'Raw ChatGPT Response',
            'step4': 'Final Restored Response',
            'input_placeholder': 'Enter your business text here...\nThe system will precisely detect company names, financial amounts, percentages, and executive names...',
            'process_btn': 'Process with Precise Detection',
            'clear_btn': 'Clear All',
            'mapping_btn': 'Show Precise Mapping Table',
            'direction': 'ltr'
        }
    else:
        return {
            'title': 'سیستم ناشناس‌سازی دقیق اطلاعات تجاری',
            'step1': 'متن ورودی و تنظیمات',
            'step2': 'متن ناشناس‌شده',
            'step3': 'پاسخ خام ChatGPT',
            'step4': 'پاسخ نهایی بازگردانده شده',
            'input_placeholder': 'متن تجاری خود را اینجا وارد کنید...\nسیستم به طور دقیق نام شرکت‌ها، مبالغ مالی، درصدها و نام مدیران را تشخیص می‌دهد...',
            'process_btn': 'پردازش با تشخیص دقیق',
            'clear_btn': 'پاک کردن همه',
            'mapping_btn': 'نمایش جدول نگاشت دقیق',
            'direction': 'rtl'
        }

def update_interface(language):
    """تغییر رابط کاربری بر اساس زبان"""
    ui_text = update_ui_text(language)
    is_english = (language == 'English')
    
    workflow_css = "workflow ltr" if is_english else "workflow rtl"
    
    return [
        gr.update(value=f"<h1 style='text-align: center; color: #FFD700; font-size: 3.5em; font-weight: bold; text-shadow: 3px 3px 6px rgba(0,0,0,0.5); margin: 20px 0; background: linear-gradient(45deg, #FFD700, #FFA500); -webkit-background-clip: text; -webkit-text-fill-color: transparent; background-clip: text;'>📊 {ui_text['title']}</h1>"),
        gr.update(value=f"<h2 style='direction: {ui_text['direction']};'>📁 {ui_text['step1']}</h2>"),
        gr.update(placeholder=ui_text['input_placeholder'], rtl=not is_english),
        gr.update(value=f"🚀 {ui_text['process_btn']}"),
        gr.update(value=f"🗑️ {ui_text['clear_btn']}"),
        gr.update(rtl=not is_english),
        gr.update(value=f"<h2 style='direction: {ui_text['direction']};'>🎭 {ui_text['step2']}</h2>"),
        gr.update(rtl=not is_english),
        gr.update(value=f"<h2 style='direction: {ui_text['direction']};'>🤖 {ui_text['step3']}</h2>"),
        gr.update(rtl=not is_english),
        gr.update(value=f"<h2 style='direction: {ui_text['direction']};'>✅ {ui_text['step4']}</h2>"),
        gr.update(rtl=not is_english),
        gr.update(value=f"📋 {ui_text['mapping_btn']}"),
        gr.update(rtl=not is_english),
        gr.update(elem_classes=workflow_css)
    ]

# ایجاد instance
anonymizer = PreciseAnonymizer()

# CSS
custom_css = """
body, .gradio-container {
    font-family: 'Segoe UI', Tahoma, Arial, sans-serif !important;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    min-height: 100vh !important;
    padding: 20px !important;
}

.rtl {
    direction: rtl !important;
    text-align: right !important;
}

.ltr {
    direction: ltr !important;
    text-align: left !important;
}

.workflow {
    display: grid !important;
    grid-template-columns: 1fr 1fr 1fr 1fr !important;
    gap: 25px !important;
    padding: 30px !important;
    align-items: start !important;
    align-content: start !important;
    grid-auto-rows: auto !important;
}

.gradio-textbox {
    border-radius: 10px !important;
    box-shadow: 0 4px 15px rgba(0,0,0,0.1) !important;
    min-height: 380px !important;
    max-height: 380px !important;
    height: 380px !important;
}

.gradio-textbox textarea {
    min-height: 350px !important;
    max-height: 350px !important;
    height: 350px !important;
    resize: vertical !important;
}

.status-box {
    background: linear-gradient(135deg, #4CAF50, #45a049) !important;
    border: 3px solid #2E7D32 !important;
    border-radius: 15px !important;
    padding: 15px !important;
    margin: 10px 0 !important;
    box-shadow: 0 8px 32px rgba(76, 175, 80, 0.3) !important;
    animation: pulse 2s infinite !important;
    min-height: 120px !important;
    max-height: 120px !important;
}

@keyframes pulse {
    0% { box-shadow: 0 8px 32px rgba(76, 175, 80, 0.3); }
    50% { box-shadow: 0 8px 40px rgba(76, 175, 80, 0.6); }
    100% { box-shadow: 0 8px 32px rgba(76, 175, 80, 0.3); }
}

.gradio-button {
    border-radius: 25px !important;
    font-weight: bold !important;
    transition: all 0.3s ease !important;
    margin: 5px 0 !important;
    min-height: 50px !important;
    max-height: 50px !important;
}

h1 {
    background: linear-gradient(45deg, #FFD700, #FFA500) !important;
    -webkit-background-clip: text !important;
    -webkit-text-fill-color: transparent !important;
    background-clip: text !important;
    min-height: 80px !important;
}
"""

# رابط کاربری Gradio
with gr.Blocks(title="📊 Precise Anonymization System", theme=gr.themes.Soft(), css=custom_css) as app:
    
    with gr.Row():
        language_selector = gr.Radio(
            choices=["فارسی", "English"],
            value="فارسی",
            label="Language / زبان",
            interactive=True
        )
    
    with gr.Column():
        title = gr.HTML("<h1 style='text-align: center; color: #FFD700; font-size: 3.5em; font-weight: bold; text-shadow: 3px 3px 6px rgba(0,0,0,0.5); margin: 20px 0; background: linear-gradient(45deg, #FFD700, #FFA500); -webkit-background-clip: text; -webkit-text-fill-color: transparent; background-clip: text;'>📊 سیستم ناشناس‌سازی دقیق اطلاعات تجاری</h1>")
        
        with gr.Row(elem_classes="workflow rtl") as workflow_row:
            with gr.Column():
                step1_title = gr.HTML('<h2 style="direction: rtl;">📁 متن ورودی و تنظیمات</h2>')
                
                input_text = gr.Textbox(
                    lines=15,
                    placeholder="متن تجاری خود را اینجا وارد کنید...\nسیستم به طور دقیق نام شرکت‌ها، مبالغ مالی، درصدها و نام مدیران را تشخیص می‌دهد...",
                    label="",
                    rtl=True
                )
                
                process_btn = gr.Button("🚀 پردازش با تشخیص دقیق", variant="primary")
                clear_btn = gr.Button("🗑️ پاک کردن همه", variant="stop")
                
                status = gr.Textbox(
                    label="وضعیت",
                    lines=4,
                    interactive=False,
                    rtl=True,
                    elem_classes=["status-box"]
                )

            with gr.Column():
                step2_title = gr.HTML('<h2 style="direction: rtl;">🎭 متن ناشناس‌شده</h2>')
                
                anonymized_output = gr.Textbox(
                    lines=15,
                    placeholder="متن ناشناس‌شده اینجا نمایش داده می‌شود...",
                    label="",
                    interactive=False,
                    rtl=True
                )

            with gr.Column():
                step3_title = gr.HTML('<h2 style="direction: rtl;">🤖 پاسخ خام ChatGPT</h2>')
                
                gpt_output = gr.Textbox(
                    lines=15,
                    placeholder="پاسخ خام ChatGPT اینجا نمایش داده می‌شود...",
                    label="",
                    interactive=False,
                    rtl=True
                )

            with gr.Column():
                step4_title = gr.HTML('<h2 style="direction: rtl;">✅ پاسخ نهایی بازگردانده شده</h2>')
                
                final_output = gr.Textbox(
                    lines=15,
                    placeholder="پاسخ نهایی اینجا نمایش داده می‌شود...",
                    label="",
                    interactive=False,
                    rtl=True
                )

        with gr.Row():
            with gr.Column():
                mapping_title = gr.HTML('<h2>🗂️ جدول نگاشت دقیق</h2>')
                mapping_btn = gr.Button("📋 نمایش جدول نگاشت دقیق")
                
                mapping_output = gr.Textbox(
                    lines=10,
                    label="جدول نگاشت اطلاعات",
                    interactive=False,
                    visible=False,
                    rtl=True
                )

    # Event handlers
    language_selector.change(
        fn=update_interface,
        inputs=[language_selector],
        outputs=[title, step1_title, input_text, process_btn, clear_btn, 
                status, step2_title, anonymized_output, step3_title, gpt_output, 
                step4_title, final_output, mapping_btn, mapping_output, workflow_row]
    )

    process_btn.click(
        fn=process_all_steps,
        inputs=[input_text, language_selector],
        outputs=[status, anonymized_output, gpt_output, final_output]
    )

    clear_btn.click(
        fn=clear_all,
        outputs=[input_text, anonymized_output, gpt_output, final_output, status]
    )

    mapping_btn.click(
        fn=get_mapping_table,
        inputs=[language_selector],
        outputs=[mapping_output]
    )

    mapping_btn.click(
        fn=lambda: gr.update(visible=True),
        outputs=[mapping_output]
    )

if __name__ == "__main__":
    app.launch()