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

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

class AnonymizerAdvanced:
    """ناشناس‌ساز پیشرفته با روش‌های متعدد"""
    
    def __init__(self, cerebras_key: str = None, gpt_key: str = None):
        self.cerebras_key = cerebras_key or os.getenv("CEREBRAS_API_KEY")
        self.gpt_key = gpt_key or os.getenv("OPENAI_API_KEY")
        self.mapping_table = {}  # {placeholder: original_text}
        self.reverse_mapping = {}  # {original_text: placeholder}
        
        logger.info("✅ Anonymizer Advanced مقداردهی شد")
    

    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 مورد نیاز است")
        
        try:
            prompt = f"""متن زیر را ناشناس کنید. قوانین:
1. اسامی اشخاص → person-01, person-02, ...
2. نام شرکت‌ها/سازمان‌ها → company-01, company-02, ...
3. مقادیر پولی → amount-01, amount-02, ...
4. درصدها → percent-01, percent-02, ...
5. فقط این توکن‌ها استفاده کنید
6. شماره‌های نسخه را درست حفظ کنید
7. اگر موجودیت تکرار شود از شماره قدیمی استفاده کنید

متن:
{text}

خروجی: فقط متن ناشناس شده"""
            
            response = 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": prompt}],
                    "max_tokens": 4096,
                    "temperature": 0.1
                },
                timeout=60
            )
            
            if response.status_code == 200:
                anonymized_text = response.json()['choices'][0]['message']['content'].strip()
                logger.info("✅ Cerebras: موفق")
                
                # استخراج mapping از متن
                self._extract_mapping_from_text(text, anonymized_text)
                return anonymized_text, self.mapping_table
            else:
                logger.error(f"❌ Cerebras Error: {response.status_code}")
                raise Exception(f"Cerebras API Error: {response.status_code}")
                
        except Exception as e:
            logger.error(f"❌ Cerebras Exception: {e}")
            raise
    

    def _extract_mapping_from_text(self, original: str, anonymized: str):
        """استخراج mapping از متن‌های اصلی و ناشناس شده"""
        # الگوهای موجودیت
        patterns = {
            'person': r'\b[ء-ي]+\s+[ء-ي]+(?:\s+[ء-ي]+)*\b',
            'company': r'(?:شرکت|بانک|سازمان|گروه|هلدینگ)\s+[ء-ي]+(?:\s+[ء-ي]+)*',
            'amount': r'\d+\s*(?:میلیارد|میلیون|هزار|تومان|ریال|دلار|تن)',
            'percent': r'\d+\s*(?:درصد|%|درصدی)',
        }
        
        for entity_type, pattern in patterns.items():
            for match in re.finditer(pattern, original):
                text_match = match.group()
                tokens = re.findall(f'{entity_type}-\\d+', anonymized)
                if tokens:
                    for token in tokens:
                        if token not in self.mapping_table:
                            self.mapping_table[token] = text_match
                            self.reverse_mapping[text_match] = token
    
    def analyze_with_gpt(self, anonymized_text: str) -> str:
        """تحلیل متن ناشناس‌شده با ChatGPT"""
        logger.info("🤖 ChatGPT تحلیل...")
        
        if not self.gpt_key:
            logger.warning("⚠️ GPT API Key نیست")
            return "❌ API Key موجود نیست"
        
        try:
            prompt = f"""متن مالی ناشناس‌شده زیر را تحلیل و خلاصه کنید:

متن:
{anonymized_text}

لطفاً:
1. خلاصه‌ای مختصر و معنادار بنویسید
2. نکات اصلی را مشخص کنید
3. تمام توکن‌های ناشناس (person-01, company-01, amount-01, percent-01) را حفظ کنید
4. تنها اطلاعات موجود در متن را بیان کنید"""
            
            response = requests.post(
                "https://api.openai.com/v1/chat/completions",
                headers={
                    "Authorization": f"Bearer {self.gpt_key}",
                    "Content-Type": "application/json"
                },
                json={
                    "model": "gpt-4o-mini",
                    "messages": [
                        {
                            "role": "system",
                            "content": "شما دستیار تحلیل متون مالی فارسی هستید. توکن‌های ناشناس را حفظ کنید."
                        },
                        {"role": "user", "content": prompt}
                    ],
                    "max_tokens": 1500,
                    "temperature": 0.7
                },
                timeout=60
            )
            
            if response.status_code == 200:
                gpt_response = response.json()['choices'][0]['message']['content']
                logger.info("✅ ChatGPT: تحلیل کامل")
                return gpt_response
            else:
                logger.error(f"❌ GPT Error: {response.status_code}")
                return "❌ خطای ChatGPT"
                
        except Exception as e:
            logger.error(f"❌ GPT Exception: {e}")
            return f"❌ خطا: {str(e)}"
    
    def restore_text(self, anonymized_text: str) -> str:
        """بازگردانی متن اصلی"""
        logger.info("🔄 بازگردانی...")
        
        restored = anonymized_text
        for placeholder, original in sorted(self.mapping_table.items()):
            restored = restored.replace(placeholder, original)
        
        logger.info("✅ بازگردانی کامل")
        return restored
    
    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, method: str = "cerebras"):
    """پردازش متن - 4 مرحله"""
    global anonymizer
    
    if not input_text.strip():
        return "", "", "", ""
    
    cerebras_key = os.getenv("CEREBRAS_API_KEY")
    gpt_key = os.getenv("OPENAI_API_KEY")
    
    if not anonymizer:
        anonymizer = AnonymizerAdvanced(cerebras_key, gpt_key)
    else:
        anonymizer.mapping_table = {}
        anonymizer.reverse_mapping = {}
    
    try:
        logger.info("=" * 70)
        logger.info(f"🚀 شروع پردازش - روش: Cerebras")
        logger.info("=" * 70)
        
        # ============================================
        # مرحله 1: ناشناس‌سازی
        # ============================================
        logger.info("📝 مرحله 1: ناشناس‌سازی...")
        
        anonymized_text, _ = anonymizer.anonymize_with_cerebras(input_text)
        
        logger.info(f"✅ ناشناس‌سازی: {len(anonymized_text)} کاراکتر")
        
        # ============================================
        # مرحله 2: تحلیل ChatGPT
        # ============================================
        logger.info("🤖 مرحله 2: تحلیل ChatGPT...")
        gpt_response = anonymizer.analyze_with_gpt(anonymized_text)
        logger.info(f"✅ تحلیل: {len(gpt_response)} کاراکتر")
        
        # ============================================
        # مرحله 3: بازگردانی پاسخ ChatGPT
        # ============================================
        logger.info("🔄 مرحله 3: بازگردانی...")
        restored_text = anonymizer.restore_text(gpt_response)
        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, gpt_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 "", "", "", "", ""

# Gradio Interface
css_rtl = """
.input-box { direction: rtl; text-align: right; }
.textbox textarea { direction: rtl; text-align: right; font-family: 'Tahoma', serif; }
"""

with gr.Blocks(title="سیستم ناشناس‌سازی متون", theme=gr.themes.Soft(), css=css_rtl) as app:
    
    gr.Markdown("# 🔐 سیستم ناشناس‌سازی متون مالی فارسی", elem_classes="input-box")
    gr.Markdown("#### استخراج و ناشناس‌سازی موجودیت‌های حساس", elem_classes="input-box")
    
    # ============================================
    # صفحه اول: دکمه‌ها (راست) + ورودی (چپ)
    # ============================================
    with gr.Row():
        # سمت راست: دکمه‌ها و تنظیمات
        with gr.Column(scale=1):
            gr.Markdown("### ⚙️ تنظیمات", elem_classes="input-box")
            
            method = gr.Radio(
                ["cerebras"],
                value="cerebras",
                label="🔧 روش پردازش",
                elem_classes="input-box"
            )
            
            gr.Markdown("---")
            
            with gr.Column():
                process_btn = gr.Button(
                    "▶️ پردازش",
                    variant="primary",
                    size="lg"
                )
                
                clear_btn = gr.Button(
                    "🗑️ پاک کردن",
                    variant="stop",
                    size="lg"
                )
        
        # سمت چپ: متن ورودی (بزرگ‌تر)
        with gr.Column(scale=3):
            input_text = gr.Textbox(
                lines=14,
                placeholder="متن مالی/خبری را وارد کنید...",
                label="📝 متن ورودی",
                elem_classes="textbox"
            )
    
    # ============================================
    # صفحه دوم: 3 باکس نتایج (وسط)
    # ============================================
    gr.Markdown("---")
    gr.Markdown("## 📊 نتایج پردازش", elem_classes="input-box")
    
    with gr.Row():
        # باکس 1: متن بازگردانی شده (راست)
        with gr.Column(scale=1):
            restored_text = gr.Textbox(
                lines=12,
                label="✅ متن بازگردانی شده",
                interactive=False,
                elem_classes="textbox"
            )
        
        # باکس 2: تحلیل ChatGPT (وسط)
        with gr.Column(scale=1):
            gpt_analysis = gr.Textbox(
                lines=12,
                label="🤖 تحلیل ChatGPT",
                interactive=False,
                elem_classes="textbox"
            )
        
        # باکس 3: متن ناشناس‌شده (چپ)
        with gr.Column(scale=1):
            anonymized_text = gr.Textbox(
                lines=12,
                label="🔒 متن ناشناس‌شده",
                interactive=False,
                elem_classes="textbox"
            )
    
    # ============================================
    # پایین: جدول نگاشت
    # ============================================
    gr.Markdown("---")
    
    mapping_table = gr.Textbox(
        lines=10,
        label="📋 جدول نگاشت",
        interactive=False,
        elem_classes="textbox"
    )
    
    # ============================================
    # Event Handlers
    # ============================================
    process_btn.click(
        fn=process,
        inputs=[input_text, method],
        outputs=[restored_text, gpt_analysis, anonymized_text, mapping_table]
    )
    
    clear_btn.click(
        fn=clear_all,
        outputs=[input_text, restored_text, gpt_analysis, anonymized_text, mapping_table]
    )

if __name__ == "__main__":
    print("=" * 70)
    print("🚀 سیستم ناشناس‌سازی متون در حال راه‌اندازی...")
    print("=" * 70)
    print("\n📋 نحوه استفاده:\n")
    print("1. CEREBRAS_API_KEY و OPENAI_API_KEY را تنظیم کنید")
    print("2. http://localhost:7860 را باز کنید")
    print("3. متن را وارد کنید")
    print("4. 'پردازش' را کلیک کنید\n")
    print("روش استفاده شده: Cerebras (Llama 3.3-70B)")
    print("=" * 70 + "\n")
    
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True
    )