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 AnonymizerCerebras:
    def __init__(self, api_key: str = None):
        self.api_key = api_key or os.getenv("CEREBRAS_API_KEY")
        self.mapping_table = {}
        self.counters = {
            'PERSON': 0, 'COMPANY': 0, 'AMOUNT': 0, 'PHONE': 0,
            'EMAIL': 0, 'ID_NUMBER': 0, 'DATE': 0, 'LOCATION': 0,
            'PERCENTAGE': 0
        }
        
        if not self.api_key:
            raise ValueError("❌ کلید API Cerebras یافت نشد!")
        
        logger.info("✅ Anonymizer مقداردهی شد")
    
    def call_cerebras(self, text: str) -> List[Dict]:
        """فراخوانی Cerebras API برای شناسایی موجودیت‌ها"""
        logger.info("🔄 فراخوانی Cerebras API...")
        
        prompt = f"""شما متخصص شناسایی اطلاعات حساس هستید.
متن زیر را تحلیل کنید و تمام اطلاعات حساسی را شناسایی کنید.

متن:
{text}

یک JSON Array برگردانید. هر عنصر دارای:
- "text": متن اطلاعات حساس
- "type": نوع (PERSON, COMPANY, AMOUNT, PHONE, EMAIL, ID_NUMBER, DATE, LOCATION, PERCENTAGE)

فقط JSON Array برگردانید!
"""
        
        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": "شما متخصص تشخیص اطلاعات حساس هستید. فقط JSON برگردانید."},
                        {"role": "user", "content": prompt}
                    ],
                    "max_tokens": 3000,
                    "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:
                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")
                return []
                
        except Exception as e:
            logger.error(f"❌ خطا Cerebras: {e}")
            return []
    
    def anonymize(self, text: str) -> Tuple[str, List]:
        """ناشناس‌سازی متن با Cerebras"""
        logger.info("🚀 شروع ناشناس‌سازی...")
        
        # تنظیف
        self.mapping_table = {}
        for key in self.counters:
            self.counters[key] = 0
        
        # دریافت موجودیت‌ها
        entities = self.call_cerebras(text)
        
        if not entities:
            logger.warning("⚠️ موجودیتی شناسایی نشد")
            return text, []
        
        # جایگزینی
        anonymized = text
        
        for entity in entities:
            entity_type = entity.get('type', 'UNKNOWN')
            entity_text = entity.get('text', '')
            
            if not entity_text or entity_type == 'UNKNOWN':
                continue
            
            idx = anonymized.find(entity_text)
            if idx == -1:
                continue
            
            if entity_type in self.counters:
                self.counters[entity_type] += 1
                token = f"[{entity_type}_{self.counters[entity_type]:03d}]"
                self.mapping_table[token] = entity_text
                
                anonymized = anonymized[:idx] + token + anonymized[idx + len(entity_text):]
                logger.info(f"✅ جایگزین: {entity_text} → {token}")
        
        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, value in sorted(self.mapping_table.items()):
            result += f"| `{token}` | {value} |\n"
        
        return result
    
    def restore(self, text: str) -> str:
        """بازگردانی اطلاعات اصلی"""
        logger.info("🔄 بازگردانی اطلاعات...")
        restored = text
        for token, value in self.mapping_table.items():
            restored = restored.replace(token, value)
        logger.info("✅ بازگردانی کامل")
        return restored


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

def process(input_text: str, api_key_cerebras: str, api_key_gpt: str) -> Tuple[str, str, str, str, str, str]:
    """
    روند کامل:
    1. ناشناس‌سازی با Cerebras
    2. ارسال به ChatGPT (حتما!)
    3. بازگردانی پاسخ ChatGPT
    """
    global anonymizer
    
    try:
        if not input_text.strip():
            return "", "", "", "", "❌ متن خالی است", ""
        
        if not api_key_gpt:
            return "", "", "", "", "❌ کلید ChatGPT الزامی است!", ""
        
        # ============================================
        # مرحله 1: مقداردهی
        # ============================================
        if not anonymizer:
            logger.info("🔧 مقداردهی Anonymizer...")
            anonymizer = AnonymizerCerebras(api_key_cerebras if api_key_cerebras else None)
        
        # ============================================
        # مرحله 2: ناشناس‌سازی
        # ============================================
        logger.info("\n" + "="*60)
        logger.info("مرحله 1️⃣: ناشناس‌سازی متن با Cerebras (Llama)")
        logger.info("="*60)
        
        anonymized_text, entities = anonymizer.anonymize(input_text)
        
        if not entities:
            status = "❌ موجودیتی شناسایی نشد"
            return input_text, "", "", "", status, ""
        
        status = f"✅ مرحله 1️⃣: ناشناس‌سازی کامل\n"
        status += f"   📍 {len(entities)} موجودیت شناسایی شد\n\n"
        
        # ============================================
        # مرحله 3: جدول نگاشت
        # ============================================
        logger.info("مرحله 2️⃣: ایجاد جدول نگاشت")
        mapping = anonymizer.get_mapping_table_str()
        status += f"✅ مرحله 2️⃣: جدول نگاشت ایجاد شد\n\n"
        
        # ============================================
        # مرحله 4: ارسال به ChatGPT (حتما!)
        # ============================================
        logger.info("="*60)
        logger.info("مرحله 3️⃣: ارسال متن ناشناس‌شده به ChatGPT")
        logger.info("="*60)
        
        prompt = f"""متن ناشناس‌شده زیر را بررسی و تحلیل کنید:

{anonymized_text}

لطفاً خلاصه‌ای مختصر و معنادار ارائه دهید."""
        
        logger.info(f"📤 ارسال به ChatGPT...\nPrompt:\n{prompt[:200]}...")
        
        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": 1000,
                    "temperature": 0.7
                },
                timeout=30
            )
            
            if gpt_response_obj.status_code == 200:
                gpt_response = gpt_response_obj.json()['choices'][0]['message']['content']
                logger.info("✅ پاسخ دریافت شد")
                status += f"✅ مرحله 3️⃣: پاسخ ChatGPT دریافت شد\n\n"
            else:
                error_text = gpt_response_obj.json().get('error', {}).get('message', gpt_response_obj.text)
                logger.error(f"❌ خطای ChatGPT: {error_text}")
                status += f"❌ مرحله 3️⃣: خطا در ChatGPT\n   📍 {error_text}\n"
                return input_text, anonymized_text, "", "", status, mapping
                
        except Exception as e:
            logger.error(f"❌ خطا در ارسال: {e}")
            status += f"❌ مرحله 3️⃣: خطا - {str(e)}\n"
            return input_text, anonymized_text, "", "", status, mapping
        
        # ============================================
        # مرحله 5: بازگردانی پاسخ ChatGPT
        # ============================================
        logger.info("="*60)
        logger.info("مرحله 4️⃣: بازگردانی پاسخ ChatGPT به حالت اصلی")
        logger.info("="*60)
        
        # ⭐⭐⭐ اینجا مهم است! ⭐⭐⭐
        # پاسخ ChatGPT را با جدول نگاشت بازگردانی کن
        # نه متن ناشناس شده را!
        restored_text = anonymizer.restore(gpt_response)
        
        logger.info(f"✅ بازگردانی کامل")
        status += f"✅ مرحله 4️⃣: بازگردانی پاسخ ChatGPT به حالت اصلی\n\n"
        status += "🎉 پردازش کامل شد!"
        
        logger.info("="*60)
        logger.info("📊 خلاصه:")
        logger.info(f"  • ورودی: {len(input_text)} کاراکتر")
        logger.info(f"  • ناشناس: {len(anonymized_text)} کاراکتر")
        logger.info(f"  • موجودیت‌ها: {len(entities)}")
        logger.info(f"  • نگاشت‌ها: {len(anonymizer.mapping_table)}")
        logger.info(f"  • پاسخ ChatGPT: {len(gpt_response)} کاراکتر")
        logger.info(f"  • نتیجه نهایی: {len(restored_text)} کاراکتر")
        logger.info("="*60)
        
        return input_text, anonymized_text, gpt_response, restored_text, status, mapping
        
    except Exception as e:
        logger.error(f"❌ خطا: {e}", exc_info=True)
        return "", "", "", "", f"❌ خطا: {str(e)}", ""

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

# رابط Gradio
with gr.Blocks(title="Anonymization + ChatGPT", theme=gr.themes.Soft()) as app:
    
    # هدر
    gr.Markdown("""
    # 📊 سیستم ناشناس‌سازی + تحلیل ChatGPT
    
    ### روند کار:
    1. **متن ورودی** → Cerebras (Llama 3.3-70B) → **ناشناس‌سازی**
    2. **متن ناشناس** → ChatGPT → **تحلیل و پاسخ**
    3. **پاسخ ChatGPT** → جدول نگاشت → **بازگردانی به حالت اصلی**
    """)
    
    # ردیف 1: تنظیمات
    with gr.Row():
        with gr.Column(scale=2):
            gr.Markdown("### ⚙️ تنظیمات API")
            api_key_cerebras = gr.Textbox(
                label="🔑 Cerebras API Key *",
                placeholder="برای ناشناس‌سازی",
                type="password",
                info="الزامی"
            )
            api_key_gpt = gr.Textbox(
                label="🔑 OpenAI API Key *",
                placeholder="برای تحلیل متن",
                type="password",
                info="الزامی"
            )
        
        with gr.Column(scale=1):
            gr.Markdown("### 🎮 کنترل‌ها")
            gr.Markdown("")
            process_btn = gr.Button("🚀 شروع پردازش", variant="primary", size="lg")
            clear_btn = gr.Button("🗑️ پاک کردن", variant="stop")
            gr.Markdown("")
            gr.Markdown("**هر دو API الزامی است**")
    
    # ردیف 2: 4 باکس در وسط
    with gr.Row():
        with gr.Column():
            gr.Markdown("### 📝 1️⃣ متن ورودی")
            input_text = gr.Textbox(
                lines=10,
                placeholder="متن خود را وارد کنید...\nمثال:\nشرکت بانک ملی از طریق مدیرعامل احمد علی\nدرآمد 600 میلیارد تومان را اعلام کرد",
                label=""
            )
        
        with gr.Column():
            gr.Markdown("### 🎭 2️⃣ متن ناشناس‌شده")
            gr.Markdown("*(برای ارسال به ChatGPT)*")
            anonymized_text = gr.Textbox(
                lines=10,
                label="",
                interactive=False
            )
        
        with gr.Column():
            gr.Markdown("### 🤖 3️⃣ پاسخ ChatGPT")
            gr.Markdown("*(قبل از بازگردانی)*")
            gpt_response = gr.Textbox(
                lines=10,
                label="",
                interactive=False
            )
        
        with gr.Column():
            gr.Markdown("### ✅ 4️⃣ نتیجه نهایی")
            gr.Markdown("*(پاسخ بازگردانی‌شده)*")
            restored_text = gr.Textbox(
                lines=10,
                label="",
                interactive=False
            )
    
    # ردیف 3: وضعیت و جدول نگاشت
    with gr.Row():
        with gr.Column():
            gr.Markdown("### ⚙️ وضعیت پردازش")
            status = gr.Textbox(
                lines=8,
                label="",
                interactive=False
            )
        
        with gr.Column():
            gr.Markdown("### 📋 جدول نگاشت")
            gr.Markdown("*(برای بازگردانی)*")
            mapping = gr.Textbox(
                lines=8,
                label="",
                interactive=False
            )
    
    # Event handlers
    process_btn.click(
        fn=process,
        inputs=[input_text, api_key_cerebras, api_key_gpt],
        outputs=[input_text, anonymized_text, gpt_response, restored_text, status, mapping]
    )
    
    clear_btn.click(
        fn=clear,
        outputs=[input_text, anonymized_text, gpt_response, restored_text, status, mapping]
    )

if __name__ == "__main__":
    print("""
╔═══════════════════════════════════════════════════════════╗
║  📊 ناشناس‌سازی + تحلیل ChatGPT                        ║
║  Cerebras Llama 3.3-70B + OpenAI                         ║
╚═══════════════════════════════════════════════════════════╝

🚀 شروع:
   pip install gradio requests
   python app_correct.py

📝 روند:
   متن ورودی
       ↓
   1️⃣ Cerebras: ناشناس‌سازی
       ↓
   2️⃣ جدول نگاشت
       ↓
   3️⃣ ChatGPT: تحلیل (حتما!)
       ↓
   4️⃣ بازگردانی: پاسخ ChatGPT

💡 کلید نکته:
   - ناشناس → ChatGPT
   - ChatGPT پاسخ → بازگردانی
   - نه ناشناس را برنمی‌گردونی!
    """)
    
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True
    )