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 ImprovedAnonymizer:
    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'همراه اول', 'company'),
                (r'گروه همراه اول', 'company'),
                
                # نام‌های مختصر شرکت‌ها
                (r'بوعلی', 'company'),  # مختصر شرکت پتروشیمی بوعلی سینا
                (r'شپنا', 'company'),    # مختصر شرکت پالایش نفت اصفهان
                (r'فولاد مبارکه', 'company'),
                
                # نام‌های در علامت نقل قول
                (r'«پارس»', 'company'),
                (r'"پارس"', 'company'),
                
                # نام اشخاص کامل
                (r'مهدی اخوان بهابادی، مدیرعامل همراه اول', 'person'),
                (r'مهدی اخوان بهابادی', 'person'),
                (r'فرج‌اله قدمی', 'person'),
                
                # مبالغ مالی - الگوهای جامع‌تر
                
                # مبالغ با ویرگول
                (r'\d+,\d{3},\d{3} میلیارد ریال', 'amount'),
                (r'\d+,\d{3} میلیارد ریال', 'amount'),
                
                # مبالغ با کلمات توضیحی
                (r'\d+ هزار و \d+ میلیارد تومان', 'amount'),
                (r'\d+ هزار و \d+ دستگاه', 'amount'),
                (r'بیش از \d+ همت', 'amount'),
                (r'حدود \d+ میلیون تومان', 'amount'),
                (r'حدود \d+ میلیون دلار', 'amount'),
                (r'نزدیک به \d+ هزار میلیارد تومان', 'amount'),
                (r'بیش از \d+ میلیارد تومان', 'amount'),
                (r'حدود \d+ میلیارد تومان', 'amount'),
                
                # مبالغ با واحد همت
                (r'\d+\.?\d* همت', 'amount'),
                (r'\d+ همت', 'amount'),
                
                # مبالغ ساده
                (r'\d+ هزار میلیارد ریالی', 'amount'),
                (r'\d+ میلیارد تومانی', 'amount'),
                (r'\d+ میلیارد تومان', 'amount'),
                (r'\d+ میلیون تومان', 'amount'),
                  (r'\d+ هزارمیلیون تومان', 'amount'),
                (r'\d+ میلیارد ریال', 'amount'),
                
                (r'\d+ ریال', 'amount'),
                
                # بازه‌های مقداری
                (r'یک تا \d+\.?\d* میلیون تن', 'amount'),
                (r'\d+ تا \d+\.?\d* میلیون تن', 'amount'),
                
                # مبالغ اعشاری
                (r'\d+\.?\d* میلیون نفر', 'amount'),
                (r'\d+\.?\d* میلیون دلار', 'amount'),
                
                # درصدها - الگوهای جامع‌تر
                (r'منفی \d+ درصد', 'percent'),
                (r'\d+\.?\d* درصدی', 'percent'),
                (r'\d+\.?\d* درصد', 'percent'),
                (r'\d+ الی \d+ درصد', 'percent'),
                (r'بیش از \d+ درصد', 'percent'),
                (r'حدود \d+ درصد', 'percent'),
            ]
            
            # مرتب‌سازی patterns بر اساس طول (طولانی‌ترین اول)
            patterns.sort(key=lambda x: len(x[0]), reverse=True)
            
            # پردازش الگوها
            for pattern, category in patterns:
                matches = list(re.finditer(pattern, anonymized, re.IGNORECASE))
                
                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]:02d}"
                        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 = "📋 **Improved 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': 'Improved 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 detect company names, financial amounts, percentages, and executive names with improved accuracy...',
            'process_btn': 'Process with Improved Detection',
            'clear_btn': 'Clear All',
            'mapping_btn': 'Show Improved 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 = ImprovedAnonymizer()

# 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="📊 Improved 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()