import gradio as gr
import re
import os
import requests
import json
import logging
from typing import Dict, List, Tuple, Optional
from llm_sender_unified import create_llm_sender

# ✅ مدل‌های موجود - به‌روزرسانی نوامبر 2024
AVAILABLE_MODELS = {
    "chatgpt": [
        # GPT-5 Series (جدیدترین)
        "gpt-5.1",           # بهترین برای کدنویسی و وظایف agentic
        "gpt-5",             # مدل reasoning قبلی
        # GPT-4 Series
        "gpt-4.1",           # هوشمندترین non-reasoning
        "gpt-4o",            # قدرتمند
        "gpt-4o-mini",       # سریع و ارزان
        "gpt-4-turbo",       # سریع‌تر از GPT-4
    ],
    "grok": [
        # Grok-4 Series (جدیدترین)
        "grok-4-fast-reasoning",      # سریع با reasoning
        "grok-4-fast-non-reasoning",  # سریع بدون reasoning
        "grok-4-0709",                # نسخه پایدار
        # Grok-3 Series
        "grok-3",                     # قدرتمند
        "grok-3-mini",                # سبک
        # Grok-2 Series
        "grok-2-vision-1212",         # با قابلیت بینایی
        "grok-2-1212",                # نسخه پایدار
        "grok-2"                      # نسخه قدیمی
    ]
}

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

class AnonymizerAdvanced:
    """ناشناس‌ساز پیشرفته با روش‌های متعدد"""
    
    def __init__(
        self, 
        cerebras_key: str = None, 
        llm_provider: str = "chatgpt", 
        llm_model: str = None,
        entities_to_anonymize: List[str] = None
    ):
        self.cerebras_key = cerebras_key or os.getenv("CEREBRAS_API_KEY")
        self.llm_provider = llm_provider
        self.llm_model = llm_model
        self.entities_to_anonymize = entities_to_anonymize or ["person", "company", "amount", "percent"]
        self.mapping_table = {}
        self.reverse_mapping = {}
        
        # ایجاد LLM sender
        self._create_llm_sender()
        
        logger.info(f"✅ Anonymizer Advanced مقداردهی شد با {llm_provider}")
    
    def _create_llm_sender(self):
        """ایجاد LLM sender مناسب"""
        try:
            # ✅ همیشه از Hugging Face Secrets استفاده کن
            if self.llm_provider == "chatgpt":
                api_key = os.getenv("OPENAI_API_KEY")
                logger.info("🔑 استفاده از OPENAI_API_KEY از Secrets")
            elif self.llm_provider == "grok":
                api_key = os.getenv("XAI_API_KEY")
                logger.info("🔑 استفاده از XAI_API_KEY از Secrets")
            else:
                api_key = None
                logger.warning("⚠️ Provider ناشناخته")
            
            # ایجاد sender
            self.llm_sender = create_llm_sender(
                provider=self.llm_provider,
                api_key=api_key,
                model=self.llm_model
            )
            
            logger.info(f"✅ LLM Sender ایجاد شد: {self.llm_provider} - {self.llm_sender.model}")
            
        except Exception as e:
            logger.error(f"❌ خطا در ایجاد LLM Sender: {e}")
            # fallback to ChatGPT
            self.llm_sender = create_llm_sender("chatgpt")

    def set_llm_provider(self, provider: str, model: str = None, entities: List[str] = None):
        """تغییر provider و مدل LLM و موجودیت‌های ناشناس‌سازی"""
        self.llm_provider = provider
        self.llm_model = model
        if entities is not None:
            self.entities_to_anonymize = entities
        self._create_llm_sender()
        logger.info(f"✅ LLM تغییر یافت به: {provider} - {model}")
        logger.info(f"✅ موجودیت‌های ناشناس‌سازی: {self.entities_to_anonymize}")

    def anonymize_with_cerebras(self, text: str) -> Tuple[str, Dict]:
        """ناشناس‌سازی با Cerebras - بر اساس موجودیت‌های انتخابی"""
        logger.info("🧠 روش Cerebras...")
        
        if not self.cerebras_key:
            logger.error("❌ Cerebras API Key موجود نیست")
            raise ValueError("Cerebras API Key مورد نیاز است")
        
        # ✅ ساخت دستورات بر اساس موجودیت‌های انتخابی
        instructions = []
        instruction_number = 1
        
        if "person" in self.entities_to_anonymize:
            instructions.append(f"{instruction_number}. اسامی اشخاص → person-01, person-02, ...")
            instruction_number += 1
        
        if "company" in self.entities_to_anonymize:
            instructions.append(f"{instruction_number}. نام شرکت‌ها/سازمان‌ها → company-01, company-02, ...")
            instruction_number += 1
        
        if "amount" in self.entities_to_anonymize:
            instructions.append(f"{instruction_number}. اعداد و ارقام و مبالغ (مثل: 50 میلیارد، 100 هزار، 25.5 میلیون، ۳۰۰ دستگاه) → amount-01, amount-02, ...")
            instruction_number += 1
        
        if "percent" in self.entities_to_anonymize:
            instructions.append(f"{instruction_number}. درصدها → percent-01, percent-02, ...")
            instruction_number += 1
        
        # اگه هیچی انتخاب نشده، متن رو همون‌طور برگردون
        if not instructions:
            logger.warning("⚠️ هیچ موجودیتی برای ناشناس‌سازی انتخاب نشده!")
            return text, {}
        
        instructions_text = "\n".join(instructions)
        instructions_text += f"\n{instruction_number}. فقط این توکن‌ها استفاده کنید"
        instructions_text += f"\n{instruction_number + 1}. شماره‌های نسخه را درست حفظ کنید"
        instructions_text += f"\n{instruction_number + 2}. اگر موجودیت تکرار شود از شماره قدیمی استفاده کنید"
        
        try:
            # مرحله 1: ناشناس‌سازی متن
            # ✅ ساخت مثال برای amount (اگر انتخاب شده)
            example_text = ""
            if "amount" in self.entities_to_anonymize:
                example_text = """
مثال:
متن اصلی: "فروش 50 میلیارد ریال در سال گذشته بود."
متن ناشناس: "فروش amount-01 در سال گذشته بود."
"""
            
            prompt1 = f"""متن زیر را ناشناس کنید. قوانین:
{instructions_text}
{example_text}
متن:
{text}

خروجی: فقط متن ناشناس شده (بدون توضیح اضافی)"""
            
            response1 = requests.post(
                "https://api.cerebras.ai/v1/chat/completions",
                headers={
                    "Authorization": f"Bearer {self.cerebras_key}",
                    "Content-Type": "application/json"
                },
                json={
                    "model": "llama-3.3-70b",
                    "messages": [{"role": "user", "content": prompt1}],
                    "max_tokens": 4096,
                    "temperature": 0.1
                },
                timeout=60
            )
            
            if response1.status_code != 200:
                logger.error(f"❌ Cerebras Error: {response1.status_code}")
                raise Exception(f"Cerebras API Error: {response1.status_code}")
            
            anonymized_text = response1.json()['choices'][0]['message']['content'].strip()
            logger.info("✅ Cerebras: ناشناس‌سازی موفق")
            
            # مرحله 2: استخراج mapping - فقط برای موجودیت‌های انتخابی
            mapping_instructions = []
            json_example = "{\n"
            
            if "person" in self.entities_to_anonymize:
                mapping_instructions.append('- برای person-XX: نام کامل شخص (مثلاً "علی احمدی")')
                json_example += '  "person-01": "متن اصلی کامل",\n'
            
            if "company" in self.entities_to_anonymize:
                mapping_instructions.append('- برای company-XX: نام کامل شرکت/سازمان (مثلاً "شرکت پتروشیمی")')
                json_example += '  "company-01": "متن اصلی کامل",\n'
            
            if "amount" in self.entities_to_anonymize:
                mapping_instructions.append('- برای amount-XX: عدد + واحد (مثلاً "80 هزار تومان" یا "50 میلیارد ریال")')
                json_example += '  "amount-01": "متن اصلی کامل با واحد",\n'
            
            if "percent" in self.entities_to_anonymize:
                mapping_instructions.append('- برای percent-XX: عدد + کلمه "درصد" (مثلاً "40 درصد" نه فقط "40")')
                json_example += '  "percent-01": "عدد + درصد",\n'
            
            json_example += "  ...\n}"
            mapping_instructions_text = "\n".join(mapping_instructions)
            
            prompt2 = f"""متن اصلی:
{text}

متن ناشناس شده:
{anonymized_text}

لطفاً یک جدول mapping برای همه توکن‌های ناشناس ایجاد کن.
برای هر توکن، متن اصلی کامل آن را مشخص کن.

**مهم:** 
{mapping_instructions_text}

خروجی را به این فرمت JSON بده (فقط JSON، بدون توضیح اضافی):
{json_example}"""
            
            response2 = requests.post(
                "https://api.cerebras.ai/v1/chat/completions",
                headers={
                    "Authorization": f"Bearer {self.cerebras_key}",
                    "Content-Type": "application/json"
                },
                json={
                    "model": "llama-3.3-70b",
                    "messages": [{"role": "user", "content": prompt2}],
                    "max_tokens": 2048,
                    "temperature": 0.1
                },
                timeout=60
            )
            
            if response2.status_code == 200:
                mapping_text = response2.json()['choices'][0]['message']['content'].strip()
                mapping_text = mapping_text.replace('```json', '').replace('```', '').strip()
                
                try:
                    self.mapping_table = json.loads(mapping_text)
                    self._fix_percent_mapping()
                    self.reverse_mapping = {v: k for k, v in self.mapping_table.items()}
                    logger.info(f"✅ Mapping استخراج شد: {len(self.mapping_table)} موجودیت")
                except json.JSONDecodeError:
                    logger.warning("⚠️ خطا در parse کردن JSON mapping - استفاده از روش fallback")
                    self._extract_mapping_from_text(text, anonymized_text)
            else:
                logger.warning("⚠️ خطا در دریافت mapping - استفاده از روش fallback")
                self._extract_mapping_from_text(text, anonymized_text)
            
            return anonymized_text, self.mapping_table
                
        except Exception as e:
            logger.error(f"❌ Cerebras Exception: {e}")
            raise
    
    def _fix_percent_mapping(self):
        """اصلاح mapping برای درصدها"""
        for token, value in self.mapping_table.items():
            value_str = str(value).strip()
            
            if token.startswith('percent-'):
                if not re.search(r'(درصد|%|درصدی)', value_str):
                    self.mapping_table[token] = f"{value_str} درصد"
                    logger.info(f"✅ اصلاح {token}: '{value_str}' → '{value_str} درصد'")
            
            elif token.startswith('amount-'):
                if not re.search(r'(میلیارد|میلیون|هزار|تومان|ریال|دلار|یورو|تن)', value_str):
                    logger.warning(f"⚠️ {token}: فقط عدد '{value_str}' - واحد مشخص نیست")

    def _extract_mapping_from_text(self, original: str, anonymized: str):
        """استخراج mapping از متن‌های اصلی و ناشناس شده - فقط برای موجودیت‌های انتخابی"""
        
        # ✅ استخراج فقط توکن‌های انتخابی
        all_tokens = []
        for entity_type in self.entities_to_anonymize:
            tokens = re.findall(f'{entity_type}-\\d+', anonymized)
            all_tokens.extend([(t, entity_type) for t in tokens])
        
        all_tokens = sorted(set(all_tokens), key=lambda x: (x[1], int(x[0].split('-')[1])))
        
        # ✅ الگوهای موجودیت - فقط برای انتخابی‌ها
        patterns = {}
        if "person" in self.entities_to_anonymize:
            patterns['person'] = r'\b[ء-ي]+\s+[ء-ي]+(?:\s+[ء-ي]+)*\b'
        if "company" in self.entities_to_anonymize:
            patterns['company'] = r'(?:شرکت|بانک|سازمان|گروه|هلدینگ)\s+[ء-ي]+(?:\s+[ء-ي]+)*'
        if "amount" in self.entities_to_anonymize:
            patterns['amount'] = r'\d+(?:\.\d+)?\s*(?:میلیارد|میلیون|هزار|تومان|ریال|دلار|یورو|تن)'
        if "percent" in self.entities_to_anonymize:
            patterns['percent'] = r'\d+(?:\.\d+)?\s*(?:درصد|%|درصدی)'
        
        original_entities = {}
        for entity_type, pattern in patterns.items():
            matches = list(re.finditer(pattern, original))
            original_entities[entity_type] = [m.group().strip() for m in matches]
        
        for token, entity_type in all_tokens:
            if entity_type in original_entities and original_entities[entity_type]:
                token_num = int(token.split('-')[1]) - 1
                
                if token_num < len(original_entities[entity_type]):
                    original_text = original_entities[entity_type][token_num]
                    self.mapping_table[token] = original_text
                    self.reverse_mapping[original_text] = token
                else:
                    original_text = original_entities[entity_type][-1]
                    if token not in self.mapping_table:
                        self.mapping_table[token] = original_text
                        self.reverse_mapping[original_text] = token
    
    def analyze_with_llm(self, anonymized_text: str, analysis_prompt: str = None) -> str:
        """استفاده از LLM یکپارچه"""
        logger.info(f"🤖 {self.llm_provider.upper()} اجرای پرامپت...")
        
        if not analysis_prompt or not analysis_prompt.strip():
            logger.info("⚠️  پرامپت خالی - بدون تحلیل")
            return "⚠️ هیچ دستور تحلیل داده نشده است"
        
        # ✅ بررسی اینکه آیا مدل GPT-4 است
        is_gpt4 = self.llm_model and any(x in self.llm_model.lower() for x in ['gpt-4', 'gpt4'])
        
        if is_gpt4:
            # ✅ پرامپت ویژه GPT-4 با مثال‌های واقعی
            logger.info("🎯 استفاده از پرامپت ویژه GPT-4")
            return self._analyze_with_gpt4_prompt(anonymized_text, analysis_prompt)
        else:
            # پرامپت عادی برای GPT-5 و Grok
            return self._analyze_with_standard_prompt(anonymized_text, analysis_prompt)
    
    def _analyze_with_gpt4_prompt(self, anonymized_text: str, analysis_prompt: str) -> str:
        """پرامپت ویژه GPT-4 با few-shot examples"""
        
        # ✅ مثال‌های واقعی Few-Shot
        few_shot_examples = """
EXAMPLE 1 - CORRECT:
Input: "company-01 فروش amount-01 داشت"
Your output should be EXACTLY: "company-01 فروش amount-01 داشت"
NOT: "company-01 فروش مبلغ amount-01 داشت"

EXAMPLE 2 - CORRECT:
Input: "amount-02 به amount-03 رسید"
Your output should be EXACTLY: "amount-02 به amount-03 رسید"
NOT: "مبلغ amount-02 به amount-03 رسید"

EXAMPLE 3 - CORRECT:
Input: "company-01 سود percent-01 داشت"
Your output should be EXACTLY: "company-01 سود percent-01 داشت"
NOT: "شرکت company-01 سود درصد percent-01 داشت"
"""
        
        # لیست توکن‌های انتخابی
        tokens_list = []
        if "person" in self.entities_to_anonymize:
            tokens_list.append("person-XX")
        if "company" in self.entities_to_anonymize:
            tokens_list.append("company-XX")
        if "amount" in self.entities_to_anonymize:
            tokens_list.append("amount-XX")
        if "percent" in self.entities_to_anonymize:
            tokens_list.append("percent-XX")
        
        tokens_str = ", ".join(tokens_list)
        
        # ✅ پرامپت انگلیسی برای GPT-4 (بهتر کار می‌کند)
        combined_text = f"""You are processing anonymized Persian/Farsi text containing placeholder tokens.

ANONYMIZED TEXT:
{anonymized_text}

USER REQUEST:
{analysis_prompt}

CRITICAL RULES:
1. Use ONLY these exact tokens: {tokens_str}
2. NEVER add words before/after tokens
3. Keep the EXACT format: amount-01 (not "مبلغ amount-01" or "amount- 01")
4. Do NOT create new tokens
5. Preserve the exact structure

{few_shot_examples}

FORBIDDEN PATTERNS - NEVER USE:
❌ "مبلغ amount-01" → ✅ Use: "amount-01"
❌ "شرکت company-01" → ✅ Use: "company-01"  
❌ "فروش به amount-02" → ✅ Use: "فروش amount-02"
❌ "درصد percent-01" → ✅ Use: "percent-01"
❌ "amount- 01" (space) → ✅ Use: "amount-01"

Now process the text following these rules EXACTLY."""
        
        try:
            # ✅ temperature خیلی پایین برای GPT-4
            logger.info(f"🌡️ Temperature: 0.05 (GPT-4 ویژه)")
            
            response = self.llm_sender.send(
                combined_text, 
                lang='en',  # انگلیسی برای GPT-4
                temperature=0.05,  # خیلی خیلی پایین
                max_tokens=2000
            )
            
            # ✅ دیباگ: نمایش خروجی خام LLM
            logger.info("=" * 60)
            logger.info("🔍 DEBUG - خروجی خام GPT-4:")
            logger.info(response[:500] + "..." if len(response) > 500 else response)
            logger.info("=" * 60)
            
            # ✅ پاکسازی قوی‌تر
            cleaned_response = self._clean_llm_response(response)
            
            # ✅ دیباگ: نمایش خروجی بعد از clean
            logger.info("=" * 60)
            logger.info("🧹 DEBUG - خروجی بعد از clean:")
            logger.info(cleaned_response[:500] + "..." if len(cleaned_response) > 500 else cleaned_response)
            logger.info("=" * 60)
            
            logger.info(f"✅ GPT-4: {len(cleaned_response)} کاراکتر")
            return cleaned_response
            
        except Exception as e:
            logger.error(f"❌ GPT-4 Exception: {e}")
            return f"❌ خطا در ارتباط با GPT-4: {str(e)}"
    
    def _analyze_with_standard_prompt(self, anonymized_text: str, analysis_prompt: str) -> str:
        """پرامپت استاندارد برای GPT-5 و Grok"""
        
        tokens_instruction = []
        examples = []
        
        if "person" in self.entities_to_anonymize:
            tokens_instruction.append("person-XX")
            examples.append("✅ صحیح: person-01 در جلسه حضور داشت\n❌ غلط: آقای person-01")
        
        if "company" in self.entities_to_anonymize:
            tokens_instruction.append("company-XX")
            examples.append("✅ صحیح: company-01 فعالیت کرد\n❌ غلط: شرکت company-01")
        
        if "amount" in self.entities_to_anonymize:
            tokens_instruction.append("amount-XX")
            examples.append("✅ صحیح: فروش amount-01 بود\n❌ غلط: فروش مبلغ amount-01")
        
        if "percent" in self.entities_to_anonymize:
            tokens_instruction.append("percent-XX")
            examples.append("✅ صحیح: رشد percent-01 داشت\n❌ غلط: رشد درصد percent-01")
        
        tokens_str = ", ".join(tokens_instruction)
        examples_str = "\n".join(examples)
        
        combined_text = f"""متن ناشناس‌سازی شده:
{anonymized_text}

دستورات:
{analysis_prompt}

⚠️ قوانین مهم:
1. فقط از کدهای ناشناس موجود استفاده کن: {tokens_str}
2. هیچ کلمه‌ای قبل یا بعد از این کدها اضافه نکن
3. کد جدید ایجاد نکن
4. ساختار دقیق متن را حفظ کن

مثال‌های صحیح و غلط:
{examples_str}"""
        
        try:
            temp_to_use = 0.2
            logger.info(f"🌡️ Temperature: {temp_to_use}")
            
            response = self.llm_sender.send(
                combined_text, 
                lang='fa',
                temperature=temp_to_use,
                max_tokens=2000
            )
            
            response = self._clean_llm_response(response)
            
            logger.info(f"✅ {self.llm_provider.upper()}: {len(response)} کاراکتر")
            return response
            
        except Exception as e:
            logger.error(f"❌ {self.llm_provider.upper()} Exception: {e}")
            return f"❌ خطا در ارتباط با {self.llm_provider.upper()}: {str(e)}"
    
    def _clean_llm_response(self, text: str) -> str:
        """پاکسازی کلمات اضافی که LLM ممکن است قبل از موجودیت‌ها اضافه کرده باشد"""
        logger.info("🧹 پاکسازی کلمات اضافی...")
        
        cleaned = text
        changes_made = 0
        
        # الگوهای کلمات اضافی برای هر نوع موجودیت
        patterns = []
        
        if "person" in self.entities_to_anonymize:
            patterns.extend([
                (r'(?:آقای|خانم|شخص|فرد)\s+(person-\d+)', r'\1'),
                (r'(person-\d+)\s+(?:نامدار|محترم|عزیز)', r'\1'),
            ])
        
        if "company" in self.entities_to_anonymize:
            patterns.extend([
                (r'(?:شرکت|سازمان|گروه|هلدینگ|بانک|موسسه)\s+(company-\d+)', r'\1'),
                (r'(company-\d+)\s+(?:محترم)', r'\1'),
            ])
        
        if "amount" in self.entities_to_anonymize:
            patterns.extend([
                # ✅ الگوهای کامل برای amount - تمام حالات ممکن
                # حالت 1: کلمات قبل از amount
                (r'(?:مبلغ|رقم|عدد|قیمت|ارزش|مقدار)\s+(amount-\d+)', r'\1'),
                (r'(?:فروش|درآمد|سود|زیان|هزینه|خرج)\s+(amount-\d+)', r'\1'),
                (r'(?:دارایی|بدهی|سرمایه|پول|وام)\s+(amount-\d+)', r'\1'),
                
                # حالت 2: حروف اضافه قبل از amount
                (r'\bبه\s+(amount-\d+)', r'\1'),
                (r'\bبا\s+(amount-\d+)', r'\1'),
                (r'\bاز\s+(amount-\d+)', r'\1'),
                (r'\bتا\s+(amount-\d+)', r'\1'),
                (r'\bدر\s+(amount-\d+)', r'\1'),
                (r'\bبرای\s+(amount-\d+)', r'\1'),
                
                # حالت 3: واحدها بعد از amount (اگر نباید باشند)
                (r'(amount-\d+)\s+(?:ریال|تومان|دلار|یورو)', r'\1'),
                (r'(amount-\d+)\s+(?:میلیون|میلیارد|هزار|تریلیون)', r'\1'),
                
                # حالت 4: ترکیبات
                (r'(?:به\s+مبلغ)\s+(amount-\d+)', r'\1'),
                (r'(?:با\s+ارزش)\s+(amount-\d+)', r'\1'),
                (r'(?:در\s+حد)\s+(amount-\d+)', r'\1'),
                
                # حالت 5: فعل + amount (بدون حرف اضافه)
                (r'(?:رسید|رسیده|می\u200cرسد)\s+(amount-\d+)', r'\1'),
                (r'(?:شد|شده|می\u200cشود)\s+(amount-\d+)', r'\1'),
                (r'(?:بود|بوده|است)\s+(amount-\d+)', r'\1'),
            ])
        
        if "percent" in self.entities_to_anonymize:
            patterns.extend([
                (r'(?:درصد|%)\s+(percent-\d+)', r'\1'),
                (r'(percent-\d+)\s+(?:درصد|درصدی|%)', r'\1'),
            ])
        
        # اعمال الگوها
        for pattern, replacement in patterns:
            new_text = re.sub(pattern, replacement, cleaned)
            if new_text != cleaned:
                count = len(re.findall(pattern, cleaned))
                changes_made += count
                cleaned = new_text
                logger.info(f"  ✅ حذف '{pattern}': {count} مورد")
        
        if changes_made > 0:
            logger.info(f"✅ {changes_made} کلمه اضافی حذف شد")
        else:
            logger.info("✅ کلمه اضافی یافت نشد")
        
        return cleaned
    
    def restore_text(self, anonymized_text: str) -> str:
        """بازگردانی متن با ترتیب بهینه برای amount"""
        logger.info("🔄 بازگردانی متن...")
        
        if not self.mapping_table:
            logger.warning("⚠️ جدول نگاشت خالی است")
            return anonymized_text
        
        logger.info(f"📋 تعداد موجودیت‌ها در mapping: {len(self.mapping_table)}")
        
        # ✅ STEP 1: normalize (hyphen یونیکد و جداسازی کلمات چسبیده)
        restored = self._normalize_tokens(anonymized_text)
        
        # ✅ STEP 2: restore قوی مخصوص amount با regex (قبل از clean!)
        # این کلیدی است - باید قبل از clean انجام شود
        logger.info("🔥 بازگردانی amount با regex...")
        amount_restored_count = 0
        for placeholder, original in self.mapping_table.items():
            if placeholder.startswith("amount-"):
                # استخراج شماره
                num = placeholder.split("-")[1]
                # الگوی regex: amount [فاصله اختیاری] - [فاصله اختیاری] شماره
                pattern = rf'amount\s*-\s*{num}'
                matches = re.findall(pattern, restored)
                if matches:
                    restored = re.sub(pattern, original, restored)
                    amount_restored_count += 1
                    logger.info(f"✅ regex: {placeholder} → {original[:30]}...")
        
        if amount_restored_count > 0:
            logger.info(f"✅ {amount_restored_count} amount با regex بازگردانی شد")
        
        # ✅ STEP 3: clean (حذف کلمات اضافی)
        # حالا که amount ها restore شدن، می‌تونیم clean کنیم
        restored = self._clean_for_restore(restored)
        
        # ✅ STEP 4: replace ساده برای بقیه (person, company, percent)
        replacements_count = 0
        for placeholder, original in sorted(self.mapping_table.items(), key=lambda x: len(x[0]), reverse=True):
            # amount ها رو قبلاً restore کردیم
            if placeholder.startswith("amount-"):
                continue
                
            if placeholder in restored:
                restored = restored.replace(placeholder, original)
                replacements_count += 1
                logger.info(f"✅ {placeholder} → {original[:30]}...")
            else:
                logger.warning(f"⚠️ {placeholder} در متن یافت نشد!")
        
        total_restored = amount_restored_count + replacements_count
        logger.info(f"✅ بازگردانی کامل - {total_restored}/{len(self.mapping_table)} جایگزین شد")
        
        # ✅ STEP 5: fallback regex برای توکن‌های باقی‌مانده
        if total_restored < len(self.mapping_table):
            logger.info("🔍 تلاش برای یافتن توکن‌های گم‌شده با regex...")
            restored = self._restore_with_regex(restored)
        
        # هشدار در صورت شکست کامل
        if total_restored == 0 and len(self.mapping_table) > 0:
            logger.error("❌ هیچ توکنی جایگزین نشد! متن ورودی احتمالاً متفاوت است.")
        
        return restored
    
    def _clean_for_restore(self, text: str) -> str:
        """پاکسازی خاص برای بازگردانی (شبیه _clean_llm_response اما سبک‌تر)"""
        logger.info("🧹 پاکسازی قبل از بازگردانی...")
        
        cleaned = text
        changes_made = 0
        
        patterns = []
        
        if "amount" in self.entities_to_anonymize:
            patterns.extend([
                (r'(?:مبلغ|رقم|عدد|قیمت|ارزش|فروش|درآمد|هزینه|سود|زیان)\s+(amount-\d+)', r'\1'),
                (r'\bبه\s+(amount-\d+)', r'\1'),
                (r'\bبا\s+(amount-\d+)', r'\1'),
                (r'\bاز\s+(amount-\d+)', r'\1'),
                (r'\bتا\s+(amount-\d+)', r'\1'),
            ])
        
        for pattern, replacement in patterns:
            new_text = re.sub(pattern, replacement, cleaned)
            if new_text != cleaned:
                changes_made += re.subn(pattern, replacement, cleaned)[1]
                cleaned = new_text
        
        if changes_made > 0:
            logger.info(f"✅ {changes_made} کلمه اضافی حذف شد")
        
        return cleaned
    
    def _restore_with_regex(self, text: str) -> str:
        """بازگردانی با استفاده از regex برای پیدا کردن توکن‌های دارای کلمات اضافی"""
        restored = text
        
        for placeholder, original in self.mapping_table.items():
            # اگر قبلاً جایگزین شده، رد شو
            if placeholder not in text:
                # الگوی regex: کلمه اضافی (اختیاری) + توکن
                # مثلاً: "فروش amount-01" یا "مبلغ amount-05"
                entity_type = placeholder.split('-')[0]
                entity_num = placeholder.split('-')[1]
                
                # الگوهای مختلف
                patterns = [
                    # کلمه فارسی + فاصله + توکن
                    rf'[ء-ي]+\s+({entity_type}-{entity_num})\b',
                    # توکن + فاصله + کلمه فارسی
                    rf'\b({entity_type}-{entity_num})\s+[ء-ي]+',
                    # فاصله اضافی داخل توکن
                    rf'\b{entity_type}\s+-\s+{entity_num}\b',
                ]
                
                for pattern in patterns:
                    matches = list(re.finditer(pattern, restored))
                    if matches:
                        logger.info(f"✅ پیدا شد با regex: {pattern}")
                        for match in matches:
                            # جایگزینی کل عبارت با فقط original
                            full_match = match.group(0)
                            # اگر توکن داخل match هست، فقط اون رو جایگزین کن
                            if placeholder in full_match:
                                restored = restored.replace(full_match, full_match.replace(placeholder, original))
                            else:
                                # اگر فرمت توکن متفاوت بود
                                restored = restored.replace(full_match, original)
                            logger.info(f"✅ regex: {placeholder} → {original[:30]}...")
                        break
        
        return restored
    
    def _normalize_tokens(self, text: str) -> str:
        """نرمال‌سازی توکن‌ها - حذف فاصله‌های اضافی و hyphen یونیکد"""
        logger.info("🧹 نرمال‌سازی توکن‌ها...")
        
        normalized = text
        changes = 0
        
        # ✅ 1. نرمال‌سازی hyphen های یونیکد برای همه موجودیت‌ها
        # این hyphen ها: ‐ ‑ ‒ – — − و hyphen معمولی -
        unicode_hyphens = r'[\u2010\u2011\u2012\u2013\u2014\u2212\-]'
        
        for entity_type in self.entities_to_anonymize:
            # تبدیل همه hyphen ها به - معمولی
            pattern = rf'{entity_type}{unicode_hyphens}(\d+)'
            replacement = rf'{entity_type}-\1'
            count = len(re.findall(pattern, normalized))
            if count > 0:
                normalized = re.sub(pattern, replacement, normalized)
                changes += count
                logger.info(f"  ✅ {entity_type}: {count} hyphen یونیکد نرمال شد")
        
        # ✅ 2. حذف فاضله‌های اضافی داخل توکن
        for entity_type in self.entities_to_anonymize:
            pattern = rf'{entity_type}\s+-\s+(\d+)'
            replacement = rf'{entity_type}-\1'
            count = len(re.findall(pattern, normalized))
            if count > 0:
                normalized = re.sub(pattern, replacement, normalized)
                changes += count
                logger.info(f"  ✅ {entity_type}: {count} فاصله اضافی حذف شد")
        
        # ✅ 3. جدا کردن توکن‌ها از کلمات فارسی چسبیده (ویژه amount)
        # مثال: amount-01در → amount-01 در
        if "amount" in self.entities_to_anonymize:
            pattern = r'(amount-\d+)([ء-ي])'
            replacement = r'\1 \2'
            before = normalized
            normalized = re.sub(pattern, replacement, normalized)
            if normalized != before:
                count = len(re.findall(pattern, before))
                changes += count
                logger.info(f"  ✅ amount: {count} کلمه چسبیده جدا شد")
        
        # ✅ 4. جدا کردن توکن‌ها از نشانه‌گذاری (ویژه amount)
        # مثال: amount-01، → amount-01 ،
        if "amount" in self.entities_to_anonymize:
            pattern = r'(amount-\d+)([،؛:.!?])'
            replacement = r'\1 \2'
            before = normalized
            normalized = re.sub(pattern, replacement, normalized)
            if normalized != before:
                count = len(re.findall(pattern, before))
                changes += count
                logger.info(f"  ✅ amount: {count} نشانه‌گذاری جدا شد")
        
        if changes > 0:
            logger.info(f"✅ مجموع {changes} تغییر نرمال‌سازی")
        
        return normalized
    
    def get_mapping_table_md(self) -> str:
        """تبدیل جدول نگاشت به Markdown"""
        if not self.mapping_table:
            return "### 📋 جدول نگاشت\n\nهیچ موجودیتی شناسایی نشد"
        
        table = "### 📋 جدول نگاشت\n\n"
        table += "| شناسه | متن اصلی |\n"
        table += "|-------|----------|\n"
        
        for token, original in sorted(self.mapping_table.items()):
            table += f"| **{token}** | {original} |\n"
        
        return table

# متغیر سراسری
anonymizer = None

def process(
    input_text: str, 
    analysis_prompt: str, 
    llm_provider: str, 
    llm_model: str,
    anonymize_all: bool,
    anonymize_person: bool,
    anonymize_company: bool,
    anonymize_amount: bool,
    anonymize_percent: bool
):
    """پردازش متن - 4 مرحله"""
    global anonymizer
    
    if not input_text.strip():
        return "", "", "", ""
    
    # ✅ ساخت لیست موجودیت‌های انتخابی
    if anonymize_all:
        entities = ["person", "company", "amount", "percent"]
    else:
        entities = []
        if anonymize_person:
            entities.append("person")
        if anonymize_company:
            entities.append("company")
        if anonymize_amount:
            entities.append("amount")
        if anonymize_percent:
            entities.append("percent")
    
    # اگه هیچی انتخاب نشده
    if not entities:
        return "", "❌ لطفاً حداقل یک موجودیت برای ناشناس‌سازی انتخاب کنید", "", ""
    
    cerebras_key = os.getenv("CEREBRAS_API_KEY")
    
    # ایجاد یا آپدیت anonymizer
    if not anonymizer:
        anonymizer = AnonymizerAdvanced(
            cerebras_key, 
            llm_provider=llm_provider,
            llm_model=llm_model,
            entities_to_anonymize=entities
        )
    else:
        anonymizer.set_llm_provider(llm_provider, llm_model, entities)
        anonymizer.mapping_table = {}
        anonymizer.reverse_mapping = {}
    
    try:
        logger.info("=" * 70)
        logger.info(f"🚀 شروع پردازش - LLM: {llm_provider} ({llm_model})")
        logger.info(f"🎯 موجودیت‌های انتخابی: {', '.join(entities)}")
        logger.info("=" * 70)
        
        # مرحله 1: ناشناس‌سازی
        logger.info("🔐 مرحله 1: ناشناس‌سازی...")
        anonymized_text, _ = anonymizer.anonymize_with_cerebras(input_text)
        logger.info(f"✅ ناشناس‌سازی: {len(anonymized_text)} کاراکتر")
        
        # ✅ دیباگ: بررسی توکن‌های موجود در متن ناشناس
        logger.info("=" * 70)
        logger.info("🔍 DEBUG - توکن‌های موجود در متن ناشناس:")
        for entity_type in entities:
            tokens_found = re.findall(f'{entity_type}-\\d+', anonymized_text)
            unique_tokens = sorted(set(tokens_found))
            logger.info(f"  {entity_type}: {unique_tokens}")
        logger.info("=" * 70)
        
        # مرحله 2: LLM (فقط اگر analysis_prompt داده شده باشد)
        has_analysis = analysis_prompt and analysis_prompt.strip()
        
        if has_analysis:
            logger.info(f"🤖 مرحله 2: {llm_provider.upper()}...")
            llm_response = anonymizer.analyze_with_llm(anonymized_text, analysis_prompt)
            logger.info(f"✅ {llm_provider.upper()}: {len(llm_response)} کاراکتر")
        else:
            logger.info("⚠️ مرحله 2: بدون تحلیل LLM (پرامپت خالی)")
            llm_response = "⚠️ هیچ دستور تحلیل داده نشده است"
        
        # مرحله 3: بازگردانی
        logger.info("🔄 مرحله 3: بازگردانی...")
        
        # ✅ اصلاح: اگر تحلیل انجام نشده، متن ناشناس اصلی رو restore کن
        if has_analysis:
            # اگر LLM تحلیل کرده، خروجی LLM رو restore کن
            restored_text = anonymizer.restore_text(llm_response)
        else:
            # اگر تحلیل نشده، متن ناشناس اصلی رو restore کن
            restored_text = anonymizer.restore_text(anonymized_text)
        
        logger.info("✅ بازگردانی کامل")
        
        # مرحله 4: جدول نگاشت
        logger.info("📋 مرحله 4: جدول نگاشت...")
        mapping_str = anonymizer.get_mapping_table_md()
        logger.info(f"✅ {len(anonymizer.mapping_table)} موجودیت")
        
        logger.info("=" * 70)
        logger.info("✅ تمام مراحل کامل!")
        logger.info("=" * 70)
        
        return restored_text, llm_response, anonymized_text, mapping_str
        
    except Exception as e:
        logger.error(f"❌ خطا: {str(e)}", exc_info=True)
        return "", f"❌ خطا: {str(e)}", "", ""

def clear_all():
    """پاک کردن همه"""
    return "", "", "", "", "", "", True, False, False, False, False

# Gradio Interface
css_rtl = """
.input-box { 
    direction: rtl; 
    text-align: right; 
}
.textbox textarea { 
    direction: rtl; 
    text-align: right; 
    font-family: 'Tahoma', serif; 
}
.thick-divider {
    border-top: 2px solid #333;
    margin: 10px 0;
}
.compact-group {
    margin: 0;
    padding: 0;
}
.compact-checkbox label {
    padding: 5px 10px !important;
    margin: 3px 0 !important;
    font-size: 0.95em !important;
}
"""

with gr.Blocks(title="سیستم ناشناس‌سازی متون", theme=gr.themes.Soft(), css=css_rtl) as app:
    
    gr.Markdown("# 🔐 پلتفرم امن چت با مدل‌های متنوع و ناشناس‌سازی داده‌ها", elem_classes="input-box")
    
    # ردیف اول: تنظیمات مدل و انتخاب موجودیت‌ها
    with gr.Row():
        # سمت راست: تنظیمات مدل
        with gr.Column(scale=1):
            with gr.Group():
                gr.Markdown("### ⚙️ تنظیمات مدل", elem_classes="input-box")
                
                llm_provider = gr.Dropdown(
                    choices=["chatgpt", "grok"],
                    value="chatgpt",
                    label="🤖 انتخاب مدل زبانی",
                    interactive=True
                )
                
                llm_model = gr.Dropdown(
                    choices=AVAILABLE_MODELS["chatgpt"],
                    value="gpt-4o-mini",
                    label="📦 انتخاب نسخه مدل",
                    interactive=True
                )
        
        # سمت چپ: انتخاب موجودیت‌ها
        with gr.Column(scale=1):
            with gr.Group():
                gr.Markdown("### 🎯 انتخاب موجودیت‌ها", elem_classes="input-box")
                
                anonymize_all = gr.Checkbox(
                    label="✅ همه موجودیت‌ها",
                    value=True,
                    elem_classes="input-box compact-checkbox"
                )
                
                anonymize_person = gr.Checkbox(
                    label="👤 اسامی اشخاص",
                    value=False,
                    elem_classes="input-box compact-checkbox"
                )
                
                anonymize_company = gr.Checkbox(
                    label="🏢 نام شرکت‌ها",
                    value=False,
                    elem_classes="input-box compact-checkbox"
                )
                
                anonymize_amount = gr.Checkbox(
                    label="💰 ارقام مالی",
                    value=False,
                    elem_classes="input-box compact-checkbox"
                )
                
                anonymize_percent = gr.Checkbox(
                    label="📊 درصدها",
                    value=False,
                    elem_classes="input-box compact-checkbox"
                )
    
    # خط جداکننده پررنگ
    gr.Markdown("---", elem_classes="thick-divider")
    
    # ردیف دوم: دستورات پردازش و متن ورودی
    with gr.Row():
        # سمت راست: دستورات پردازش
        with gr.Column(scale=1):
            gr.Markdown("### 📋 دستورات پردازش", elem_classes="input-box")
            
            analysis_prompt = gr.Textbox(
                lines=22,
                placeholder="مثال: این متن را خلاصه کن\nیا: نکات کلیدی را استخراج کن",
                label="📋 دستورات LLM (اختیاری)",
                elem_classes="textbox"
            )
        
        # سمت چپ: متن ورودی
        with gr.Column(scale=1):
            gr.Markdown("### 📝 متن ورودی", elem_classes="input-box")
            
            input_text = gr.Textbox(
                lines=22,
                placeholder="متن مالی/خبری را وارد کنید...",
                label="",
                elem_classes="textbox"
            )
    
    # دکمه‌های پردازش و پاک کردن
    with gr.Row():
        process_btn = gr.Button(
            "▶️ پردازش",
            variant="primary",
            size="lg",
            scale=2
        )
        
        clear_btn = gr.Button(
            "🗑️ پاک کردن",
            variant="stop",
            size="lg",
            scale=1
        )
    
    # نتایج
    gr.Markdown("## 📊 نتایج پردازش", elem_classes="input-box")
    
    with gr.Row():
        with gr.Column(scale=1):
            restored_text = gr.Textbox(
                lines=12,
                label="✅ متن بازگردانی شده",
                interactive=False,
                elem_classes="textbox"
            )
        
        with gr.Column(scale=1):
            llm_analysis = gr.Textbox(
                lines=12,
                label="🤖 تحلیل LLM",
                interactive=False,
                elem_classes="textbox"
            )
        
        with gr.Column(scale=1):
            anonymized_text = gr.Textbox(
                lines=12,
                label="🔒 متن ناشناس‌شده",
                interactive=False,
                elem_classes="textbox"
            )
    
    mapping_table = gr.Markdown(
        value="### 📋 جدول نگاشت\n\nهنوز پردازشی انجام نشده",
        label="📋 جدول نگاشت",
        elem_classes="input-box"
    )
    
    
    # Event Handler برای تغییر provider
    def handle_provider_change(provider):
        models = AVAILABLE_MODELS.get(provider, [])
        default_model = models[0] if models else None
        return gr.update(choices=models, value=default_model)
    
    llm_provider.change(
        fn=handle_provider_change,
        inputs=[llm_provider],
        outputs=[llm_model]
    )
    
    def handle_select_all(select_all):
        if select_all:
            return (
                gr.update(value=False, interactive=False),
                gr.update(value=False, interactive=False),
                gr.update(value=False, interactive=False),
                gr.update(value=False, interactive=False)
            )
        else:
            return (
                gr.update(value=False, interactive=True),
                gr.update(value=False, interactive=True),
                gr.update(value=False, interactive=True),
                gr.update(value=False, interactive=True)
            )
    
    anonymize_all.change(
        fn=handle_select_all,
        inputs=[anonymize_all],
        outputs=[anonymize_person, anonymize_company, anonymize_amount, anonymize_percent]
    )
    
    # پردازش
    process_btn.click(
        fn=process,
        inputs=[
            input_text, 
            analysis_prompt, 
            llm_provider, 
            llm_model,
            anonymize_all,
            anonymize_person,
            anonymize_company,
            anonymize_amount,
            anonymize_percent
        ],
        outputs=[restored_text, llm_analysis, anonymized_text, mapping_table]
    )
    
    # پاک کردن
    clear_btn.click(
        fn=clear_all,
        outputs=[
            input_text, 
            analysis_prompt,
            restored_text, 
            llm_analysis, 
            anonymized_text, 
            mapping_table,
            anonymize_all,
            anonymize_person,
            anonymize_company,
            anonymize_amount,
            anonymize_percent
        ]
    )

if __name__ == "__main__":
    print("=" * 70)
    print("🚀 سیستم ناشناس‌سازی متون در حال راه‌اندازی...")
    print("=" * 70)
    print("\n📋 نحوه استفاده:\n")
    print("1. API Keyها را در Hugging Face Secrets تنظیم کنید:")
    print("   - CEREBRAS_API_KEY (ضروری برای ناشناس‌سازی)")
    print("   - OPENAI_API_KEY (برای ChatGPT)")
    print("   - XAI_API_KEY (برای Grok)")
    print("2. http://localhost:7860 را باز کنید")
    print("3. مدل زبانی (ChatGPT/Grok) و نسخه مدل را انتخاب کنید")
    print("4. موجودیت‌های مورد نظر برای ناشناس‌سازی را انتخاب کنید")
    print("5. متن و دستورات پردازش را وارد کنید")
    print("6. 'پردازش' را کلیک کنید\n")
    print("🔐 تمام API Keyها از Hugging Face Secrets خوانده می‌شوند")
    print("📦 مدل‌های پشتیبانی شده:")
    print("   • ChatGPT GPT-5: gpt-5.1, gpt-5")
    print("   • ChatGPT GPT-4: gpt-4.1, gpt-4o, gpt-4o-mini, gpt-4-turbo")
    print("   • Grok-4: grok-4-fast-reasoning, grok-4-fast-non-reasoning, grok-4-0709")
    print("   • Grok-3: grok-3, grok-3-mini")
    print("   • Grok-2: grok-2-vision-1212, grok-2-1212, grok-2")
    print("=" * 70 + "\n")
    
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True
    )