import requests
import json
import gradio as gr
from typing import Dict, Any
import os
from dataclasses import dataclass
import re

@dataclass
class CerebrasConfig:
    """تنظیمات Cerebras API برای Qwen 3-32B"""
    api_key: str
    base_url: str = "https://api.cerebras.ai/v1"
    model: str = "qwen/qwen-3-14b:free"
    max_tokens: int = 16384
    temperature: float = 0.3        # ⬇️ کمتر = دقیق‌تر
    top_p: float = 0.8              # ⬇️ کمتر = محافظه‌کارانه‌تر
    presence_penalty: float = 0.2   # ⬆️ بیشتر = کمتر تکرار
    frequency_penalty: float = 0.2  # ⬆️ بیشتر = تنوع بیشتر
    #model: str = "qwen-3-32b"
   #max_tokens: int = 3000  # افزایش برای thinking tokens
    #temperature: float = 0.3  # کاهش برای دقت بیشتر
   #top_p: float = 0.9

class AdvancedCerebrasAnonymizer:
    """سیستم پیشرفته ناشناس‌سازی متون مالی/خبری فارسی با Qwen 3-32B"""
    
    def __init__(self, api_key: str = None):
        if api_key is None:
            api_key = os.getenv("CEREBRAS_API_KEY")
            if not api_key:
                raise ValueError("کلید API یافت نشد")
        
        self.config = CerebrasConfig(api_key=api_key)
        self.system_prompt = self._create_advanced_system_prompt()
    
    def _create_advanced_system_prompt(self) -> str:
        """ایجاد دستورالعمل سیستمی بهینه شده"""
        return """شما یک سیستم ناشناس‌سازی متون مالی فارسی هستید.

⚠️ CRITICAL: در پاسخ نهایی خود، فقط و فقط متن ناشناس‌سازی شده را برگردانید، بدون هیچ توضیح، تحلیل، یا تگ اضافی.

## قوانین اندیس‌گذاری:
1. **ترتیب پیوسته**: company-01, company-02, ... | person-01, person-02, ... | amount-01, amount-02, ... | percent-01, percent-02, ...
2. **ثبات**: اگر "همراه اول" → company-01 شد، در تمام متن همان باشد
3. **نام مستعار**: "فاما" = "فولاد مبارکه" → هر دو company-01
4. **اشاره ضمنی**: "این شرکت" اگر به company-01 اشاره دارد → company-01 (نه company-02)

## انواع موجودیت:
- **company-XX**: شرکت‌ها، بانک‌ها، سازمان‌ها، گروه‌ها
- **person-XX**: نام و نام خانوادگی اشخاص  
- **amount-XX**: مبالغ - واحد را حفظ کن
- **percent-XX**: درصدها

## قوانین کلیدی:
1. بازرس = شرکت است → company-XX
2. واحدها: "amount-01 میلیارد تومان" ✅
3. گروه‌ها: "گروه X" → company-XX
4. کلمات عمومی حفظ: "سه شرکت" → حفظ
5. دوره زمانی حفظ: "۵ ماهه" → حفظ
6. بازه = یک entity: "یک تا 1.5 میلیون" → amount-01
7.شناسایی و دسته‌بندی درصدها بین 50 تا 70 درصد به عنوان یک موجودیت درصد در متن
8.شناسایی تمام ارقام موجود در متن به عنوان موجودیت amount-XX مانند "سود خالص 50 میلیارد تومان" را به مقدار amount-01 تبدیل کن.
9.شناسایی مقدار درصد در بازه 40–60٪ به عنوان یک موجودیت درصد (percent-03مثلا).
10.بزرگ‌ترین هلدینگ شستا در حوزه بازار سرمایه عمل می‌کند" را به صورت "بزرگ‌ترین هلدینگ company-03 در حوزه بازار سرمایه عمل می‌کند"
11.هر جا که یک شرکت و گروه با نام یکسان ذکر شود، آن را به یک موجودیت یکسان (company-XX) تبدیل کن.
12.سود حاصل از منابع عملیاتی ۱۰،۸۸۷،۸۶۴ میلیون ریال" را به عنوان amount-01 شناسایی کن.
13."بانک ملی ایران" را به عنوان company-01 شناسایی کن.
14.شرکت ارتباطات سیار ایران همراه اول" را به عنوان company-01 شناسایی کن.
15."سپرده‌گذاری مرکزی اوراق بهادار و تسویه وجوه" را به عنوان company-01 شناسایی کن.

## مثال:
ورودی: ایران خودرو در اسفند 1402 حدود 23 هزار و 296 میلیارد درآمد کسب کرد که 4.58 درصد افزایش داشت.
خروجی: company-01 در اسفند 1402 حدود amount-01 درآمد کسب کرد که percent-01 افزایش داشت.
"این اپراتور در منطقه توانسته به پوششی 100 میلیونی دست پیدا کند" را به صورت "این اپراتور در منطقه توانسته به پوششی amount-04 دست پیدا کند".


⚠️ یادآوری: فقط متن ناشناس‌شده، بدون هیچ توضیح اضافی."""

    def _make_api_request(self, text: str) -> Dict[str, Any]:
        """ارسال درخواست به Cerebras API"""
        headers = {
            "Authorization": f"Bearer {self.config.api_key}",
            "Content-Type": "application/json"
        }
        
        # اضافه کردن /no_think به متن برای غیرفعال کردن reasoning
        user_content = f"{text}\n\n/no_think"
        
        payload = {
            "messages": [
                {
                    "role": "system",
                    "content": self.system_prompt
                },
                {
                    "role": "user", 
                    "content": user_content
                }
            ],
            "model": self.config.model,
            "temperature": self.config.temperature,
            "top_p": self.config.top_p,
            "max_tokens": self.config.max_tokens
        }
        
        try:
            response = requests.post(
                f"{self.config.base_url}/chat/completions",
                headers=headers,
                json=payload,
                timeout=60
            )
            response.raise_for_status()
            return response.json()
            
        except requests.exceptions.RequestException as e:
            raise Exception(f"خطا در ارتباط با Cerebras API: {str(e)}")
    
    def anonymize_text(self, text: str) -> Dict[str, Any]:
        """ناشناس‌سازی متن با استفاده از Qwen 3-32B"""
        if not text.strip():
            return {
                "success": False,
                "error": "متن ورودی خالی است"
            }
        
        try:
            response = self._make_api_request(text)
            
            if "choices" not in response or not response["choices"]:
                return {
                    "success": False,
                    "error": "پاسخ نامعتبر از API"
                }
            
            content = response["choices"][0]["message"]["content"]
            
            # پاک کردن thinking tags اگر وجود دارد
            content = self._remove_thinking_tags(content)
            
            # پاک کردن markdown اگر وجود دارد
            content = self._clean_markdown(content)
            
            # حذف خطوط اضافی و فضاهای خالی
            content = content.strip()
            
            # حذف توضیحات اضافی در ابتدا یا انتها
            content = self._clean_explanations(content)
            
            # تحلیل نتایج
            analysis = self._analyze_anonymized_text(content)
            
            return {
                "success": True,
                "anonymized_text": content,
                "entities": analysis["entities"],
                "statistics": analysis["statistics"],
                "detailed_analysis": analysis["detailed_analysis"],
                "usage": response.get("usage", {}),
                "quality_check": self._validate_anonymized_text(content)
            }
                
        except Exception as e:
            return {
                "success": False,
                "error": f"خطا در پردازش: {str(e)}"
            }
    
    def _remove_thinking_tags(self, content: str) -> str:
        """حذف تگ‌های thinking از خروجی"""
        # حذف محتوای داخل <think>...</think>
        content = re.sub(r'<think>.*?</think>', '', content, flags=re.DOTALL)
        
        # حذف تگ‌های خالی
        content = re.sub(r'</?think>', '', content)
        
        return content.strip()
    
    def _clean_markdown(self, content: str) -> str:
        """پاک کردن markdown از پاسخ"""
        if "```" in content:
            lines = content.split('\n')
            clean_lines = []
            skip = False
            for line in lines:
                if line.strip().startswith('```'):
                    skip = not skip
                    continue
                if not skip:
                    clean_lines.append(line)
            content = '\n'.join(clean_lines)
        return content
    
    def _clean_explanations(self, content: str) -> str:
        """حذف توضیحات اضافی در ابتدا یا انتها"""
        lines = content.split('\n')
        clean_lines = []
        
        for line in lines:
            # حذف خطوطی که شامل توضیحات متا هستند
            if any(word in line.lower() for word in ['okay', 'let me', 'here is', 'خروجی', 'نتیجه', 'پاسخ']):
                continue
            clean_lines.append(line)
        
        return '\n'.join(clean_lines).strip()
    
    def _analyze_anonymized_text(self, text: str) -> Dict[str, Any]:
        """تحلیل متن ناشناس‌سازی شده"""
        import re
        
        # شمارش موجودیت‌ها
        companies = re.findall(r'company-(\d+)', text)
        persons = re.findall(r'person-(\d+)', text)
        amounts = re.findall(r'amount-(\d+)', text)
        percents = re.findall(r'percent-(\d+)', text)
        
        # آمار کلی
        statistics = {
            "company": len(set(companies)),
            "person": len(set(persons)),
            "amount": len(set(amounts)),
            "percent": len(set(percents)),
            "total_replacements": len(companies) + len(persons) + len(amounts) + len(percents)
        }
        
        # جزئیات موجودیت‌ها
        entities = {
            "companies": sorted(list(set(companies)), key=lambda x: int(x)),
            "persons": sorted(list(set(persons)), key=lambda x: int(x)),
            "amounts": sorted(list(set(amounts)), key=lambda x: int(x)),
            "percents": sorted(list(set(percents)), key=lambda x: int(x))
        }
        
        # تحلیل دقیق‌تر
        detailed_analysis = {
            "preserved_dates": len(re.findall(r'\d{4}/\d{1,2}/\d{1,2}|\d{1,2}\s+\w+\s+\d{4}', text)),
            "preserved_times": len(re.findall(r'\d{1,2}:\d{2}', text)),
            "financial_indicators": len(re.findall(r'\b(EPS|P/E|ARPU|NPL|ROE|ROA)\b', text)),
            "units_preserved": len(re.findall(r'(میلیارد|میلیون|هزار|تومان|ریال|درهم|دلار|یورو|تن|کیلوگرم)', text))
        }
        
        return {
            "statistics": statistics,
            "entities": entities,
            "detailed_analysis": detailed_analysis
        }
    
    def _validate_anonymized_text(self, text: str) -> Dict[str, Any]:
        """اعتبارسنجی پیشرفته متن ناشناس‌شده"""
        import re
        
        # استخراج همه موجودیت‌ها
        companies = re.findall(r'company-(\d+)', text)
        persons = re.findall(r'person-(\d+)', text)
        amounts = re.findall(r'amount-(\d+)', text)
        percents = re.findall(r'percent-(\d+)', text)
        
        validation_issues = []
        
        # بررسی هر نوع موجودیت
        for entity_type, indices in [
            ("company", companies),
            ("person", persons),
            ("amount", amounts),
            ("percent", percents)
        ]:
            if indices:
                unique_indices = sorted(list(set([int(x) for x in indices])))
                # بررسی شروع از 1
                if unique_indices[0] != 1:
                    validation_issues.append(f"اندیس {entity_type} از 01 شروع نشده (شروع: {unique_indices[0]:02d})")
                
                # بررسی پیوستگی
                expected = list(range(1, len(unique_indices) + 1))
                if unique_indices != expected:
                    validation_issues.append(f"اندیس‌های {entity_type} پیوسته نیستند: {[f'{x:02d}' for x in unique_indices]}")
        
        # بررسی کلمات انگلیسی غیرضروری
        english_words = re.findall(r'\b[a-zA-Z]+\b', text)
        unwanted_english = [word for word in english_words if word.lower() not in ['eps', 'p/e', 'arpu', 'npl', 'roe', 'roa']]
        if unwanted_english:
            validation_issues.append(f"کلمات انگلیسی غیرضروری: {unwanted_english}")
        
        return {
            "is_valid": len(validation_issues) == 0,
            "issues": validation_issues,
            "entity_counts": {
                "company": len(set(companies)),
                "person": len(set(persons)),
                "amount": len(set(amounts)),
                "percent": len(set(percents))
            }
        }

# باقی کد رابط کاربری همان قبل است...
# (تابع create_advanced_interface و بقیه کد تغییری ندارد)

def create_advanced_interface():
    """ایجاد رابط کاربری پیشرفته"""
    
    api_key_available = bool(os.getenv("CEREBRAS_API_KEY"))
    
    custom_css = """
    .gradio-container {
        font-family: 'Tahoma', 'Arial', sans-serif !important;
        direction: rtl;
        max-width: 1400px;
        margin: 0 auto;
    }
    .result-box {
        background-color: #f8f9fa;
        border: 2px solid #e9ecef;
        border-radius: 12px;
        padding: 20px;
        margin: 10px 0;
    }
    .warning-box {
        background-color: #fff3cd;
        border: 2px solid #ffeaa7;
        border-radius: 12px;
        padding: 15px;
        color: #856404;
        margin: 10px 0;
    }
    .success-box {
        background-color: #d4edda;
        border: 2px solid #c3e6cb;
        border-radius: 12px;
        padding: 15px;
        color: #155724;
        margin: 10px 0;
    }
    .qwen-box {
        background-color: #e7f3ff;
        border: 2px solid #2196F3;
        border-radius: 12px;
        padding: 15px;
        color: #0d47a1;
        margin: 10px 0;
    }
    .stats-grid {
        display: grid;
        grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
        gap: 15px;
        margin: 15px 0;
    }
    .stat-card {
        background-color: #ffffff;
        border: 1px solid #dee2e6;
        border-radius: 8px;
        padding: 15px;
        text-align: center;
        box-shadow: 0 2px 4px rgba(0,0,0,0.1);
    }
    .quality-badge {
        display: inline-block;
        padding: 5px 10px;
        border-radius: 20px;
        font-weight: bold;
        margin: 5px;
    }
    .quality-pass {
        background-color: #28a745;
        color: white;
    }
    .quality-fail {
        background-color: #dc3545;
        color: white;
    }
    """
    
    with gr.Blocks(css=custom_css, title="ناشناس‌ساز پیشرفته با Qwen 3-32B", theme=gr.themes.Soft()) as interface:
        
        gr.Markdown("""
        # 🔒 سیستم پیشرفته ناشناس‌سازی متون مالی/خبری فارسی
        ### ⚡ قدرت‌گرفته از Cerebras AI + Alibaba Qwen 3-32B
        """)
        
        gr.Markdown("""
        <div class="qwen-box">
        🚀 <strong>مدل: Alibaba Qwen 3-32B (اصلاح شده)</strong><br>
        ⚡ سرعت: 2,400 توکن/ثانیه | 🧠 32B پارامتر | 💰 $0.40/$0.80<br>
        ✅ <strong>بهینه‌سازی شده:</strong> Thinking Mode غیرفعال برای خروجی مستقیم
        </div>
        """)
        
        if api_key_available:
            gr.Markdown("""
            <div class="success-box">
            ✅ <strong>سیستم آماده است</strong> - کلید API تنظیم شده
            </div>
            """)
            api_key_input = gr.Textbox(visible=False, value="")
        else:
            gr.Markdown("""
            <div class="warning-box">
            ⚠️ <strong>کلید API تنظیم نشده</strong><br>
            لطفاً کلید Cerebras API خود را وارد کنید
            </div>
            """)
            api_key_input = gr.Textbox(
                label="🔑 کلید Cerebras API",
                placeholder="csk-...",
                type="password"
            )
        
        with gr.Row():
            with gr.Column(scale=1):
                input_text = gr.Textbox(
                    label="📝 متن ورودی",
                    placeholder="متن مالی یا خبری خود را اینجا وارد کنید...",
                    lines=12,
                    max_lines=25
                )
                
                with gr.Row():
                    anonymize_btn = gr.Button(
                        "🔒 ناشناس‌سازی", 
                        variant="primary", 
                        size="lg"
                    )
                    clear_btn = gr.Button(
                        "🗑️ پاک کردن", 
                        variant="secondary"
                    )
            
            with gr.Column(scale=1):
                output_text = gr.Textbox(
                    label="🎯 متن ناشناس‌سازی شده",
                    lines=12,
                    max_lines=25,
                    elem_classes=["result-box"]
                )
                
                copy_btn = gr.Button(
                    "📋 کپی متن",
                    variant="secondary",
                    size="sm"
                )
        
        copy_output = gr.Textbox(
            label="📋 متن برای کپی",
            lines=3,
            visible=False,
            interactive=True
        )
        
        with gr.Row():
            with gr.Column():
                statistics_output = gr.Markdown(label="📊 آمار کلی")
            with gr.Column():
                quality_output = gr.Markdown(label="✅ کنترل کیفیت")
        
        with gr.Row():
            with gr.Column():
                entities_output = gr.Markdown(label="🏷️ موجودیت‌ها")
            with gr.Column():
                detailed_analysis_output = gr.Markdown(label="🔍 تحلیل دقیق")
        
        usage_output = gr.Markdown(label="⚡ اطلاعات پردازش")
        
        def process_advanced_text(text: str, api_key_manual: str = ""):
            """پردازش متن"""
            if api_key_manual is None:
                api_key_manual = ""
            
            final_api_key = ""
            if api_key_manual and api_key_manual.strip():
                final_api_key = api_key_manual.strip()
            elif os.getenv("CEREBRAS_API_KEY"):
                final_api_key = os.getenv("CEREBRAS_API_KEY")
            
            if not final_api_key:
                return ("", "❌ کلید API وارد نشده است", "", "", "", "")
            
            if not text or not text.strip():
                return ("", "❌ لطفاً متن ورودی را وارد کنید", "", "", "", "")
            
            try:
                anonymizer = AdvancedCerebrasAnonymizer(api_key=final_api_key)
                result = anonymizer.anonymize_text(text)
                
                if not result["success"]:
                    return ("", f"❌ خطا: {result['error']}", "", "", "", "")
                
                stats = result.get("statistics", {})
                stats_md = f"""📊 **آمار کلی:**

<div class="stats-grid">
    <div class="stat-card"><h3>🏢 شرکت‌ها</h3><h2>{stats.get('company', 0)}</h2></div>
    <div class="stat-card"><h3>👤 اشخاص</h3><h2>{stats.get('person', 0)}</h2></div>
    <div class="stat-card"><h3>💰 مبالغ</h3><h2>{stats.get('amount', 0)}</h2></div>
    <div class="stat-card"><h3>📊 درصدها</h3><h2>{stats.get('percent', 0)}</h2></div>
    <div class="stat-card"><h3>🔢 کل</h3><h2>{stats.get('total_replacements', 0)}</h2></div>
</div>
"""
                
                quality = result.get("quality_check", {})
                quality_md = "✅ **کنترل کیفیت:**\n\n"
                
                if quality.get("is_valid", False):
                    quality_md += '<span class="quality-badge quality-pass">✅ موفق</span>\n\n'
                else:
                    quality_md += '<span class="quality-badge quality-fail">❌ مشکل</span>\n\n'
                    issues = quality.get("issues", [])
                    if issues:
                        quality_md += "**مشکلات:**\n"
                        for issue in issues:
                            quality_md += f"• {issue}\n"
                
                entities = result.get("entities", {})
                entities_md = "🏷️ **موجودیت‌ها:**\n\n"
                
                if entities.get("companies"):
                    entities_md += f"🏢 company-{', company-'.join(entities['companies'])}\n\n"
                if entities.get("persons"):
                    entities_md += f"👤 person-{', person-'.join(entities['persons'])}\n\n"
                if entities.get("amounts"):
                    entities_md += f"💰 amount-{', amount-'.join(entities['amounts'])}\n\n"
                if entities.get("percents"):
                    entities_md += f"📊 percent-{', percent-'.join(entities['percents'])}\n\n"
                
                detailed = result.get("detailed_analysis", {})
                detailed_md = f"""🔍 **تحلیل:**

📅 تاریخ: {detailed.get('preserved_dates', 0)}
📈 شاخص‌ها: {detailed.get('financial_indicators', 0)}
📏 واحدها: {detailed.get('units_preserved', 0)}
"""
                
                usage = result.get("usage", {})
                usage_md = f"""⚡ **Qwen 3-32B:**

🤖 مدل: {anonymizer.config.model}
🌡️ Temperature: {anonymizer.config.temperature}
📥 Input: {usage.get('prompt_tokens', '?')}
📤 Output: {usage.get('completion_tokens', '?')}
📊 Total: {usage.get('total_tokens', '?')}
"""
                
                return (
                    result["anonymized_text"],
                    stats_md,
                    quality_md,
                    entities_md,
                    detailed_md,
                    usage_md
                )
                
            except Exception as e:
                return ("", f"❌ خطا: {str(e)}", "", "", "", "")
        
        def copy_text(text_to_copy):
            if not text_to_copy or not text_to_copy.strip():
                return gr.Textbox(visible=False), "⚠️ متن خالی"
            return gr.Textbox(value=text_to_copy, visible=True), "✅ آماده کپی"
        
        def clear_all():
            return "", "", "", "", "", "", "", gr.Textbox(visible=False)
        
        anonymize_btn.click(
            fn=process_advanced_text,
            inputs=[input_text, api_key_input],
            outputs=[output_text, statistics_output, quality_output, entities_output, detailed_analysis_output, usage_output]
        )
        
        copy_btn.click(
            fn=copy_text,
            inputs=[output_text],
            outputs=[copy_output, statistics_output]
        )
        
        clear_btn.click(
            fn=clear_all,
            outputs=[input_text, output_text, statistics_output, quality_output, entities_output, detailed_analysis_output, usage_output, copy_output]
        )
        
        gr.Examples(
            examples=[
                ["ایران خودرو در اسفندماه سال 1402 حدود 23 هزار و 296 میلیارد تومان درآمد کسب کرد که در مقایسه با بهمن 4.58 درصد افزایش داشت."],
                ["مجمع پتروشیمی بوعلی سینا برگزار شد. وانیا نیک تدبیر را بازرس قانونی انتخاب کردند."],
                ["شرکت فولاد مبارکه اصفهان با ملی نفت قرارداد امضا کرد. فاما سرمایه را از 8700 به 12500 میلیارد افزایش می‌دهد."]
            ],
            inputs=input_text,
            label="📚 مثال‌ها"
        )
        
        with gr.Accordion("📖 راهنما", open=False):
            gr.Markdown("""
            ## 🔧 اصلاحات انجام شده:
            
            ### ✅ مشکل حل شده:
            - غیرفعال کردن Thinking Mode با `/no_think`
            - حذف خودکار تگ‌های `<think>`
            - پاک‌سازی توضیحات اضافی
            - کاهش Temperature به 0.3 برای دقت بیشتر
            
            ### 🎯 چگونه کار می‌کند:
            1. پرامپت بهینه شده برای خروجی مستقیم
            2. اضافه کردن `/no_think` به انتهای درخواست
            3. پردازش و پاک‌سازی خودکار خروجی
            4. حذف تمام توضیحات و تگ‌های اضافی
            """)
        
        return interface

if __name__ == "__main__":
    interface = create_advanced_interface()
    interface.launch(
        show_error=True,
        ssr_mode=False  # برای پایداری بیشتر
    )