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

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

class AnonymizerCerebrasFixed:
    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
        }
        
        if not self.api_key:
            raise ValueError("❌ کلید API Cerebras یافت نشد!")
        
        logger.info("✅ Anonymizer Fixed مقداردهی شد")
    
    def get_system_prompt(self) -> str:
        """System Prompt برای ناشناس‌سازی مستقیم"""
        return """شما یک «ناشناس‌ساز متون مالی فارسی پیشرفته» هستید.

## وظایف:
1. تمام اسامی خاص را ناشناس کنید
2. تمام مقادیر عددی را ناشناس کنید
3. Entity Linking دقیق اعمال کنید

## قوانین CRITICAL:

### برای اشخاص:
- "سروش خسروی" و "خسروی" و "سروش" و "وی" = person-01
- نام خانوادگی تنهایی را رها نکنید
- ضمیرهای اشاره را ناشناس کنید

### برای شرکت‌ها:
- "پتروشیمی غدیر" و "غدیر" و "این شرکت" = company-01
- نام کوتاه را ناشناس کنید

### برای مقادیر:
- واحد یکسان = ID یکسان
- "142 میلیارد" و "۱۴۲ میلیارد" = amount-01

### برای درصدها:
- "21 درصد" و "21%" و "۲۱ درصدی" = percent-01

## ترتیب شماره‌گذاری:
- company-01, company-02, ...
- person-01, person-02, ...
- amount-01, amount-02, ...
- percent-01, percent-02, ...

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

{text}

دستورات:
1. نام کامل + نام تنهایی + ضمیرها = یک موجودیت
2. واحد یکسان = موجودیت یکسان
3. Entity Linking دقیق اعمال کنید
4. جدول نگاشت بسازید

خروجی: متن ناشناس‌شده"""
    
    def anonymize_with_api(self, text: str) -> Tuple[str, Dict]:
        """ناشناس‌سازی با استفاده از API"""
        logger.info("🚀 شروع ناشناس‌سازی...")
        
        system_prompt = self.get_system_prompt()
        user_prompt = self.get_user_prompt(text)
        
        try:
            logger.info("📡 ارسال به Cerebras API...")
            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": 4096,
                    "temperature": 0.1
                },
                timeout=60
            )
            
            if response.status_code != 200:
                error_msg = response.text
                logger.error(f"❌ خطای API: {error_msg}")
                return text, {"error": error_msg}
            
            result = response.json()
            anonymized_text = result['choices'][0]['message']['content'].strip()
            
            logger.info(f"✅ ناشناس‌سازی کامل شد")
            logger.info(f"متن ورودی: {len(text)} کاراکتر")
            logger.info(f"متن خروجی: {len(anonymized_text)} کاراکتر")
            
            return anonymized_text, {"status": "success"}
            
        except Exception as e:
            logger.error(f"❌ خطا: {str(e)}")
            return text, {"error": str(e)}
    
    def anonymize_with_regex(self, text: str) -> Tuple[str, Dict]:
        """ناشناس‌سازی ساده با Regex (بدون API)"""
        logger.info("🔍 شروع ناشناس‌سازی Regex...")
        
        anonymized = text
        mapping = {}
        entity_count = {'person': 0, 'company': 0, 'amount': 0, 'percent': 0}
        
        # الگوهای Regex
        patterns = {
            'person': r'\b[ء-ي]+\s+[ء-ي]+(?:\s+[ء-ي]+)*\b',  # نام‌های فارسی
            'amount': r'\d+\s*(?:میلیارد|میلیون|هزار|تومان|ریال|دلار|تن|دستگاه)',
            'percent': r'\d+\s*(?:درصد|%|درصدی)',
            'company': r'(?:شرکت|بانک|سازمان|گروه|هلدینگ)\s+[ء-ي]+(?:\s+[ء-ي]+)*',
        }
        
        # ناشناس‌سازی
        for entity_type, pattern in patterns.items():
            matches = re.finditer(pattern, anonymized)
            for match in matches:
                text_match = match.group()
                key = text_match.lower()
                
                if key not in mapping:
                    entity_count[entity_type] += 1
                    placeholder = f"{entity_type}-{entity_count[entity_type]:02d}"
                    mapping[key] = placeholder
                    anonymized = anonymized.replace(text_match, placeholder, 1)
        
        logger.info(f"✅ {sum(entity_count.values())} موجودیت ناشناس شد")
        
        return anonymized, mapping
    
    def get_mapping_table_str(self) -> str:
        """تبدیل جدول نگاشت به string"""
        if not self.mapping_table:
            return "### جدول نگاشت خالی است"
        
        table = "### 📋 جدول نگاشت\n\n"
        table += "| ID | نوع | متن اصلی |\n"
        table += "|----|----|----------|\n"
        
        for token, info in self.mapping_table.items():
            entity_type = info.get('type', 'نامشخص')
            original = info.get('original', '')
            table += f"| {token} | {entity_type} | {original} |\n"
        
        return table

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

def process(input_text: str, api_choice: str = "cerebras"):
    """پردازش متن"""
    global anonymizer
    
    if not input_text.strip():
        return "", "", ""
    
    # مقداردهی
    api_key = os.getenv("CEREBRAS_API_KEY")
    
    if not api_key and api_choice == "cerebras":
        return "", "", "❌ API Key تنظیم نشده است"
    
    if not anonymizer:
        if api_key:
            anonymizer = AnonymizerCerebrasFixed(api_key)
        else:
            anonymizer = AnonymizerCerebrasFixed("dummy")  # Regex mode
    
    try:
        if api_choice == "cerebras" and api_key:
            logger.info("استفاده از Cerebras API")
            anonymized_text, result = anonymizer.anonymize_with_api(input_text)
        else:
            logger.info("استفاده از Regex")
            anonymized_text, mapping = anonymizer.anonymize_with_regex(input_text)
            anonymizer.mapping_table = {v: {'original': k, 'type': 'unknown'} for k, v in mapping.items()}
        
        mapping_str = anonymizer.get_mapping_table_str()
        
        return input_text, anonymized_text, mapping_str
        
    except Exception as e:
        logger.error(f"❌ خطا: {str(e)}")
        return "", "", f"❌ خطا: {str(e)}"

def clear():
    """پاک کردن"""
    return "", "", ""

# Gradio Interface
css_rtl = """
#input_text textarea { direction: rtl; text-align: right; }
#anonymized_text textarea { direction: rtl; text-align: right; }
#mapping 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():
        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>")
            
            api_choice = gr.Radio(
                ["cerebras", "regex"],
                value="regex",
                label="انتخاب روش"
            )
            
            process_btn = gr.Button("🔄 پردازش", variant="primary", size="lg")
            clear_btn = gr.Button("🗑️ پاک کردن", variant="stop", size="lg")
    
    with gr.Row():
        anonymized_text = gr.Textbox(
            lines=12,
            label="🔒 متن ناشناس‌شده",
            interactive=False,
            elem_id="anonymized_text"
        )
    
    with gr.Row():
        mapping = gr.Textbox(
            lines=8,
            label="📋 جدول نگاشت",
            interactive=False,
            elem_id="mapping"
        )
    
    # Event handlers
    process_btn.click(
        fn=process,
        inputs=[input_text, api_choice],
        outputs=[input_text, anonymized_text, mapping]
    )
    
    clear_btn.click(
        fn=clear,
        outputs=[input_text, anonymized_text, mapping]
    )

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