import gradio as gr
import re
import os
import requests
import logging
from typing import Dict, List, Tuple, Set
import json

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class AnonymizerCerebrasEnhanced:
    def __init__(self, api_key: str = None):
        self.api_key = api_key or os.getenv("CEREBRAS_API_KEY")
        self.mapping_table = {}
        self.counters = {
            'company': 0, 'person': 0, 'amount': 0, 'phone': 0,
            'email': 0, 'id_number': 0, 'date': 0, 'location': 0,
            'percent': 0
        }
        self.seen_entities = {}  # برای ثبات نگاشت
        
        if not self.api_key:
            raise ValueError("❌ کلید API Cerebras یافت نشد!")
        
        logger.info("✅ Anonymizer Enhanced مقداردهی شد")
    
    def get_system_prompt(self) -> str:
        """ایجاد دستورالعمل سیستمی پیشرفته برای Groq"""
        return """شما یک «ناشناس‌ساز متون مالی/خبری فارسی» هستید. وظیفه‌تان جایگزینی اسامی خاص و مقادیر عددی با شناسه‌های بی‌معناست.

## **قوانین اندیس‌گذاری - CRITICAL**
### **1. ترتیب شماره‌گذاری الزامی:**
- شرکت‌ها: company-01, company-02, company-03, company-04, ... (پیوسته و بدون گپ)
- اشخاص: person-01, person-02, person-03, ... (پیوسته و بدون گپ)
- اعداد: amount-01, amount-02, amount-03, ... (پیوسته و بدون گپ)
- درصدها: percent-01, percent-02, percent-03, ... (پیوسته و بدون گپ)
- تاریخها: date-01, date-02, date-03, ... (پیوسته و بدون گپ)

### **2. ثبات شناسه‌ها در متن - MUST MAINTAIN:**
- اگر "همراه اول" اول‌بار company-01 شد، در تمام متن همان باشد
- اگر "مهدی احمدی" اول‌بار person-01 شد، در تمام متن همان باشد
- **CRITICAL: اگر "سروش خسروی" = person-01، تو "خسروی" تنهایی هم = person-01 باشد**

### **3. تشخیص صحیح انواع:**
**شرکت/سازمان:** همراه اول، بانک ملی، ایران‌خودرو، سایپا، بانک مرکزی، سامانه کدال، وزارت نفت
**شخص:** مهدی اخوان بهابادی، محمدرضا فرزین، ابوالفضل نجارزاده، سروش خسروی
**عدد:** 37، 70، 677، 73.7، 178 (هر عددی)
**درصد:** 37 درصدی، 15 درصدی، 53 درصد، 43%
**تاریخ:** 1403، 1404، اردیبهشت، فروردین، 30 آذر 1403

## **انواع موجودیت‌ها:**

**company-XX:** نام شرکت‌ها، سازمان‌ها، بانک‌ها، هلدینگ‌ها، گروه‌های مالی

**person-XX:** نام و نام خانوادگی اشخاص - شامل نام کامل، نام کوچک تنهایی، نام خانوادگی تنهایی

**amount-XX:** مبالغ مالی شامل ریال، تومان، همت، دلار، تن، دستگاه و واحدهای اندازه‌گیری

**percent-XX:** درصدها و نسبت‌ها

**date-XX:** تمام تاریخ‌ها شامل سال، ماہ، روز و ترکیب آنها

## **قوانین کلیدی:**

1. **ترتیب شماره‌گذاری:** اولین باری که موجودیت ظاهر می‌شود، شماره می‌گیرد (01، 02، 03، ...)

2. **حفظ هویت یکسان:** اگر همان موجودیت دوباره آمد، از همان شماره استفاده کن.

3. **CRITICAL - Entity Linking برای اشخاص:**
   - اگر "سروش خسروی" = person-01 شد، تو "خسروی" تنهایی = person-01
   - اگر "محمدرضا فرزین" = person-01 شد، تو "فرزین" یا "محمدرضا" = person-01
   - اگر "علی احمدی" = person-01 شد، تو "احمدی"، "علی"، "آن شخص" همه = person-01
   - **MUST TRACK: نام کامل → نام کوچک → نام خانوادگی → ضمیرها**
   - **نام خانوادگی تنهایی را هرگز بدون linking رها نکن**

4. **تشخیص نام‌های مختلف:** "فولاد مبارکه اصفهان" و "فولاد مبارکه" و "این شرکت" همه company-01 هستند.

5. **CRITICAL - تمام تاریخ‌ها باید Anonymize شوند:**
   - سال ONLY: "سال 1403" → "سال date-01"
   - ماہ ONLY: "اردیبهشت" → "date-02"
   - سال + ماہ: "اردیبهشت 1404" → "date-03 date-04"
   - تاریخ مکمل: "1403/04/12" → "date-05/date-06/date-07"
   - **NO EXCEPTION: تمام اعداد تاریخ باید anonymize شوند**
   - **یکسانی برقرار کن: اگر "1403" یک جا date-01 شد، همه جا date-01 باشد**

6. **مبالغ و درصدهای مختلف:** هر عدد جدید، شماره جدید می‌گیرد

7. **حفظ ساختار:** ساختار جمله را حفظ کن، کلمات توصیفی مثل "شرکت"، "بانک"، "گروه" را قبل از برچسب حفظ کن

8. **هیچ توضیح اضافه‌ای نده:** فقط متن ناشناس‌شده را برگردان

## **موارد حفظ شده:**
- عناوین شغلی: مدیرعامل، رئیس کل، مدیرکل، سرپرست
- واحدها: میلیارد تومان، همت، ریال، ماه، سال
- مکان‌ها: تهران، اصفهان، ایران
- کلمات توضیحی: "شرکت"، "بانک"، "گروه"

## **ممنوع:**
- کلمات انگلیسی اضافی
- تغییر ساختار جمله
- حذف یا اضافه کردن کلمات
- **نام خانوادگی یا نام کوچک تنهایی را بدون linking رها کردن**

## **نمونه‌های آموزشی:**

**نمونه ۱ - Entity Linking برای نام‌ها (CRITICAL):**
ورودی: سروش خسروی، سرپرست هیأت‌مدیره. خسروی اعلام کرد که سود خالص 216 میلیارد تومان بود. خسروی همچنین به چالش‌ها اشاره کرد.
خروجی: person-01، سرپرست هیأت‌مدیره. person-01 اعلام کرد که سود خالص amount-01 بود. person-01 همچنین به چالش‌ها اشاره کرد.

**نمونه ۲ - تمام تاریخها Anonymize شوند (CRITICAL):**
ورودی: سال 1403 یکی از سخت‌ترین سال‌ها برای صنعت پتروشیمی بود. در اردیبهشت 1404 شرکت گزارش منتشر کرد و کاهش سود خالص در سال 1403 را اعلام کرد.
خروجی: سال date-01 یکی از سخت‌ترین سال‌ها برای صنعت پتروشیمی بود. در date-02 date-03 شرکت گزارش منتشر کرد و کاهش سود خالص در سال date-01 را اعلام کرد.

**نمونه ۳:**
ورودی: مهدی اخوان بهابادی، مدیرعامل همراه اول، اعلام کرد درآمد عملیاتی شرکت با رشد 37 درصدی به 70 هزار و 677 میلیارد تومان رسیده است. سود خالص 7101 میلیارد تومان و تلفیقی گروه همراه اول 8003 میلیارد تومان شد.
خروجی: person-01، مدیرعامل company-01، اعلام کرد درآمد عملیاتی شرکت با رشد percent-01 به amount-01 رسیده است. سود خالص amount-02 و تلفیقی گروه company-01 amount-03 شد.

**نمونه ۴:**
ورودی: بانک مرکزی و بانک ملی با همکاری محمدرضا فرزین، 60 درصد سپرده‌ها را مدیریت کردند.
خروجی: company-01 و company-02 با همکاری person-01، percent-01 سپرده‌ها را مدیریت کردند.

**نمونه ۵:**
ورودی: سایپا و ایران‌خودرو مجموع زیان 620 همت داشتند و سایپا 269 هزار میلیارد زیان اعلام کرد.
خروجی: company-01 و company-02 مجموع زیان amount-01 داشتند و company-01 amount-02 زیان اعلام کرد.

**نمونه ۶ - تاریخ مکمل:**
ورودی: مجمع عمومی مورخ 1403/04/12 برگزار شد و گزارش مالی منتهی به 30 آذر 1403 تصویب رسید.
خروجی: مجمع عمومی مورخ date-01/date-02/date-03 برگزار شد و گزارش مالی منتهی به date-04 date-05 date-06 تصویب رسید.

**نمونه ۷:**
ورودی: بانک پاسارگاد با شناسایی سود خالص 155 هزار میلیارد ریالی در رده دوم سودآورترین بانک‌های کشور قرار گرفت. در مقابل، بانک سرمایه با مدیرعاملی فرج‌اله قدمی وضعیت بحرانی دارد.
خروجی: company-01 با شناسایی سود خالص amount-01 در رده دوم سودآورترین بانک‌های کشور قرار گرفت. در مقابل، company-02 با مدیرعاملی person-01 وضعیت بحرانی دارد.

**نمونه ۸:**
ورودی: بانک سرمایه با مدیرعاملی فرج‌اله قدمی زیان خالص 2700 میلیارد تومانی در سه‌ماهه نخست 1404 گزارش کرد. نسبت کفایت سرمایه به منفی 345 درصد رسیده و زیان انباشته نزدیک به 67 هزار میلیارد تومان است.
خروجی: company-01 با مدیرعاملی person-01 زیان خالص amount-01 در سه‌ماهه نخست date-01 گزارش کرد. نسبت کفایت سرمایه به منفی percent-01 رسیده و زیان انباشته نزدیک به amount-02 است.

**نمونه ۹:**
ورودی: دو بانک ملت و پاسارگاد به ترتیب با شناسایی سود خالص 157 و 155 هزار میلیارد ریالی رقابت تنگاتنگی داشته و در رده‌های اول و دوم جای دارند.
خروجی: دو بانک company-01 و company-02 به ترتیب با شناسایی سود خالص amount-01 و amount-02 رقابت تنگاتنگی داشته و در رده‌های اول و دوم جای دارند.

**نمونه ۱۰:**
ورودی: مرور صورت‌های مالی بانک‌ها نشان می‌دهد سهم سودهای ارزی به‌راحتی به 40–60٪ رسیده است و این مسئله نشان‌دهنده وضعیت غیرعادی بازار است.
خروجی: مرور صورت‌های مالی بانک‌ها نشان می‌دهد سهم سودهای ارزی به‌راحتی به percent-01 رسیده است و این مسئله نشان‌دهنده وضعیت غیرعادی بازار است.

**فقط متن ناشناس‌شده را برگردان - هیچ توضیح اضافی نیاز نیست."""
    
    def get_user_prompt(self, text: str) -> str:
        """تشکیل پرامپت کاربر"""
        return f"""متن مالی فارسی زیر را تجزیه و تحلیل کنید. تمام موجودیت‌های حساس را شناسایی کنید و یک JSON Array برگردانید.

متن:
{text}

**مهم**: 
- اگر چند بار یک نام تکرار شود، یک id بدهید
- کلمات عمومی را حفظ کنید
- واحدها را حفظ کنید
- فقط JSON برگردانید!

یک JSON Array برگردانید. هر عنصر دارای:
- "text": متن دقیق استخراج شده
- "type": نوع (company, person, amount, percent, phone, email, date, location, id_number)
- "original": توضیح اضافی اگر نام مستعار باشد"""

    def call_cerebras(self, text: str) -> List[Dict]:
        """فراخوانی Cerebras API با پرامپت بهبود شده"""
        logger.info("🔄 فراخوانی Cerebras API با دستورالعمل قوی...")
        
        system_prompt = self.get_system_prompt()
        user_prompt = self.get_user_prompt(text)
        
        try:
            response = requests.post(
                "https://api.cerebras.ai/v1/chat/completions",
                headers={
                    "Authorization": f"Bearer {self.api_key}",
                    "Content-Type": "application/json"
                },
                json={
                    "model": "llama-3.3-70b",
                    "messages": [
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": user_prompt}
                    ],
                    "max_tokens": 4000,
                    "temperature": 0.1
                },
                timeout=30
            )
            
            if response.status_code != 200:
                logger.error(f"❌ خطای API Cerebras: {response.text}")
                return []
            
            result = response.json()
            content = result['choices'][0]['message']['content']
            
            try:
                # تمیز کردن محتوا از markdown اگر وجود داشته باشد
                content = content.replace("```json", "").replace("```", "").strip()
                entities = json.loads(content)
                if not isinstance(entities, list):
                    entities = []
                logger.info(f"✅ {len(entities)} موجودیت استخراج شد")
                return entities
            except json.JSONDecodeError:
                logger.error(f"❌ خطا در JSON parsing: {content[:200]}")
                return []
                
        except Exception as e:
            logger.error(f"❌ خطا Cerebras: {e}")
            return []
    
    def get_placeholder(self, entity_type: str) -> str:
        """تولید placeholder با format جدید"""
        type_lower = entity_type.lower()
        if type_lower not in self.counters:
            type_lower = 'amount'
        
        self.counters[type_lower] += 1
        return f"{type_lower}-{self.counters[type_lower]:02d}"
    
    def anonymize(self, text: str) -> Tuple[str, List]:
        """ناشناس‌سازی متن با قوانین ثبات"""
        logger.info("🚀 شروع ناشناس‌سازی متن...")
        
        # تنظیف
        self.mapping_table = {}
        self.seen_entities = {}
        for key in self.counters:
            self.counters[key] = 0
        
        # دریافت موجودیت‌ها
        entities = self.call_cerebras(text)
        
        if not entities:
            logger.warning("⚠️ موجودیتی شناسایی نشد")
            return text, []
        
        logger.info("🔄 Processing entities...")
        
        # جایگزینی با قانون ثبات
        anonymized = text
        replacements = []
        
        for entity in entities:
            entity_type = entity.get('type', 'amount').lower()
            entity_text = entity.get('text', '').strip()
            original_info = entity.get('original', '')
            
            if not entity_text:
                continue
            
            # بررسی اگر این موجودیت قبلاً دیده شده است
            entity_key = (entity_type, entity_text.lower())
            
            if entity_key in self.seen_entities:
                token = self.seen_entities[entity_key]
                logger.info(f"🔄 موجودیت تکراری: {entity_text} → {token}")
            else:
                token = self.get_placeholder(entity_type)
                self.seen_entities[entity_key] = token
                self.mapping_table[token] = {
                    'original': entity_text,
                    'type': entity_type,
                    'note': original_info
                }
                logger.info(f"✅ جایگزینی: {entity_text} → {token}")
            
            # جایگزینی دقیق (case-sensitive اول، سپس case-insensitive)
            idx = anonymized.find(entity_text)
            if idx != -1:
                anonymized = anonymized[:idx] + token + anonymized[idx + len(entity_text):]
                replacements.append({
                    'original': entity_text,
                    'placeholder': token,
                    'type': entity_type
                })
        
        logger.info(f"✅ ناشناس‌سازی کامل - {len(self.mapping_table)} نگاشت")
        return anonymized, entities
    
    def get_mapping_table_str(self) -> str:
        """جدول نگاشت جزئی"""
        if not self.mapping_table:
            return "❌ موجودیتی شناسایی نشد"
        
        result = "## 📊 جدول نگاشت\n\n"
        result += "| توکن | اطلاعات اصلی | نوع |\n"
        result += "|------|--------|------|\n"
        
        for token, info in sorted(self.mapping_table.items()):
            entity_type = info.get('type', 'unknown')
            original = info.get('original', '')
            note = info.get('note', '')
            
            note_str = f" ({note})" if note else ""
            result += f"| `{token}` | {original}{note_str} | {entity_type} |\n"
        
        return result
    
    def restore(self, text: str) -> str:
        """بازگردانی اطلاعات اصلی"""
        logger.info("🔄 بازگردانی اطلاعات...")
        restored = text
        for token, info in self.mapping_table.items():
            original = info.get('original', '')
            restored = restored.replace(token, original)
        logger.info("✅ بازگردانی کامل")
        return restored


# متغیرهای global
anonymizer = None

def process(input_text: str) -> Tuple[str, str, str, str, str]:
    """
    روند کامل:
    1. ناشناس‌سازی با Cerebras (llama-3.3-70b) + پرامپت قوی
    2. ارسال به ChatGPT (حتما!)
    3. بازگردانی پاسخ ChatGPT
    """
    global anonymizer
    
    try:
        if not input_text.strip():
            return "", "", "", "", ""
        
        # دریافت API Keys
        api_key_cerebras = os.getenv("CEREBRAS_API_KEY")
        api_key_gpt = os.getenv("OPENAI_API_KEY")
        
        if not api_key_gpt:
            logger.error("❌ OPENAI_API_KEY یافت نشد")
            return "", "", "", "", ""
        
        if not api_key_cerebras:
            logger.error("❌ CEREBRAS_API_KEY یافت نشد")
            return "", "", "", "", ""
        
        # ============================================
        # مرحله 1: مقداردهی
        # ============================================
        if not anonymizer:
            logger.info("Initializing anonymizer...")
            anonymizer = AnonymizerCerebrasEnhanced()
        
        # ============================================
        # مرحله 2: ناشناس‌سازی با پرامپت قوی
        # ============================================
        logger.info("Step 1: Anonymizing text with Cerebras...")
        
        anonymized_text, entities = anonymizer.anonymize(input_text)
        
        if not entities:
            logger.warning("⚠️ موجودیتی شناسایی نشد - متن ناشناس نشد")
            return input_text, "", "", "", ""
        
        # ============================================
        # مرحله 3: جدول نگاشت
        # ============================================
        logger.info("Step 2: Creating mapping table")
        mapping = anonymizer.get_mapping_table_str()
        logger.info(f"📋 {len(anonymizer.mapping_table)} نگاشت ایجاد شد")
        
        # ============================================
        # مرحله 4: ارسال به ChatGPT (حتما!)
        # ============================================
        logger.info("Step 3: Sending to ChatGPT...")
        
        prompt = f"""متن ناشناس‌شده زیر (متن مالی) را تحلیل و خلاصه کنید.

متن:
{anonymized_text}

لطفاً:
1. خلاصه‌ای مختصر و معنادار ارائه دهید
2. نکات اصلی را مشخص کنید
3. تمام توکن‌های ناشناس (مثل company-01، amount-02) را حفظ کنید
4. تنها اطلاعات موجود در متن را بیان کنید"""
        
        logger.info(f"📤 ارسال به ChatGPT (gpt-4o-mini)...")
        
        try:
            gpt_response_obj = requests.post(
                "https://api.openai.com/v1/chat/completions",
                headers={"Authorization": f"Bearer {api_key_gpt}"},
                json={
                    "model": "gpt-4o-mini",
                    "messages": [
                        {
                            "role": "system", 
                            "content": "شما دستیار تحلیل متون مالی فارسی هستید. متن‌های ناشناس‌شده را دقیق تحلیل کنید. تمام توکن‌های ناشناس را حفظ کنید."
                        },
                        {"role": "user", "content": prompt}
                    ],
                    "max_tokens": 1500,
                    "temperature": 0.7
                },
                timeout=30
            )
            
            if gpt_response_obj.status_code == 200:
                gpt_response = gpt_response_obj.json()['choices'][0]['message']['content']
                logger.info("✅ پاسخ دریافت شد")
            else:
                error_text = gpt_response_obj.json().get('error', {}).get('message', gpt_response_obj.text)
                logger.error(f"❌ خطای ChatGPT: {error_text}")
                return input_text, anonymized_text, "", "", mapping
                
        except Exception as e:
            logger.error(f"❌ خطا در ارسال به ChatGPT: {e}")
            return input_text, anonymized_text, "", "", mapping
        
        # ============================================
        # مرحله 5: بازگردانی پاسخ ChatGPT
        # ============================================
        logger.info("Step 4: Restoring original text...")
        
        restored_text = anonymizer.restore(gpt_response)
        
        logger.info(f"✅ بازگردانی کامل")
        
        logger.info(f"Done. Input: {len(input_text)} | Anonymized: {len(anonymized_text)} | Entities: {len(entities)}")
        
        return input_text, anonymized_text, gpt_response, restored_text, mapping
        
    except Exception as e:
        logger.error(f"❌ خطا عمومی: {e}", exc_info=True)
        return "", "", "", "", ""

def clear():
    """پاک کردن"""
    empty_mapping = "### 📋 جدول نگاشت\nدر انتظار پردازش..."
    return "", "", "", "", empty_mapping

# رابط Gradio - کاملاً فارسی‌زبان و RTL
css_rtl = """
#input_text textarea { direction: rtl; text-align: right; }
#anonymized_text textarea { direction: rtl; text-align: right; }
#gpt_response textarea { direction: rtl; text-align: right; }
#restored_text textarea { direction: rtl; text-align: right; }
"""

with gr.Blocks(title="سیستم ناشناس‌سازی متون", theme=gr.themes.Soft(), css=css_rtl) as app:
    
    gr.Markdown("# 🔐 سیستم ناشناس‌سازی متون مالی فارسی")
    gr.Markdown("#### استخراج موجودیت‌های حساس و ناشناس‌سازی آنها")
    
    with gr.Row():
        # بلوک 1: متن ورودی (سمت راست)
        with gr.Column(scale=2):
            input_text = gr.Textbox(
                lines=12,
                placeholder="متن مالی/خبری را وارد کنید...",
                label="📝 متن ورودی",
                elem_id="input_text"
            )
        
        # دکمه‌های کنترل
        with gr.Column(scale=1):
            gr.HTML("<div style='text-align: center; margin-bottom: 10px;'></div>")
            process_btn = gr.Button("🔄 پردازش", variant="primary", size="lg")
            clear_btn = gr.Button("🗑️ پاک کردن", variant="stop", size="lg")
    
    # بلوک 2: متن ناشناس‌سازی شده
    with gr.Row():
        with gr.Column(scale=1):
            anonymized_text = gr.Textbox(
                lines=10,
                label="🔒 متن ناشناس‌شده",
                interactive=False,
                elem_id="anonymized_text"
            )
        
        # بلوک 3: پاسخ ChatGPT
        with gr.Column(scale=1):
            gpt_response = gr.Textbox(
                lines=10,
                label="🤖 تحلیل ChatGPT",
                interactive=False,
                elem_id="gpt_response"
            )
        
        # بلوک 4: متن بازگردانی شده (سمت چپ)
        with gr.Column(scale=1):
            restored_text = gr.Textbox(
                lines=10,
                label="✅ متن بازگردانی شده",
                interactive=False,
                elem_id="restored_text"
            )
    
    # بلوک 5: جدول نگاشت به صورت مارکداون
    with gr.Row():
        with gr.Column():
            mapping = gr.Markdown(
                value="### 📋 جدول نگاشت\nدر انتظار پردازش...",
                label="📋 جدول نگاشت"
            )
    
    # Event handlers
    process_btn.click(
        fn=process,
        inputs=[input_text],
        outputs=[input_text, anonymized_text, gpt_response, restored_text, mapping]
    )
    
    clear_btn.click(
        fn=clear,
        outputs=[input_text, anonymized_text, gpt_response, restored_text, mapping]
    )

if __name__ == "__main__":
    print("🚀 سیستم ناشناس‌سازی متون در حال راه‌اندازی...")
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True
    )