# --- START OF FILE app.py ---

import os
import json
import logging
import threading
import base64
import io
import time
from flask import Flask, render_template, request, Response, stream_with_context
import requests
import docx

# ================== تنظیمات لاگ ==================
# فقط خطاهای حیاتی را لاگ میکنیم تا کنسول شلوغ نشود
logging.basicConfig(level=logging.ERROR, format='%(asctime)s - %(levelname)s - %(message)s')
app = Flask(__name__)

# ================== تنظیمات کلیدها ==================

GEMINI_MODEL_NAME = "gemini-2.5-flash" 
ALL_KEYS_STR = os.environ.get("ALL_GEMINI_API_KEYS", "")
# تمیز کردن کلیدها
GEMINI_API_KEYS = [key.strip() for key in ALL_KEYS_STR.split(',') if key.strip()]

key_index_counter = 0
key_lock = threading.Lock()

def get_next_key_with_index():
    global key_index_counter
    with key_lock:
        if not GEMINI_API_KEYS:
            return None, -1
        current_index = key_index_counter
        key = GEMINI_API_KEYS[current_index]
        key_index_counter = (key_index_counter + 1) % len(GEMINI_API_KEYS)
        return key, current_index

# ================== مسیرها ==================

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/chat', methods=['POST'])
def chat():
    # اگر کلا کلیدی نباشد، باز هم نباید کرش کنیم
    if not GEMINI_API_KEYS:
        # شبیه سازی پاسخ متنی ربات به جای ارور سیستمی
        fake_response = {"choices": [{"delta": {"content": "⚠️ تنظیمات سیستم کامل نیست (کلید API یافت نشد)."}}]}
        return Response(f"data: {json.dumps(fake_response)}\n\n", mimetype='text/event-stream')

    data = request.json
    system_instruction = "تو چت بات هوش مصنوعی آلفا هستی. دوستانه، کوتاه و با ایموجی پاسخ بده."
    show_thoughts = data.get("show_thoughts", False)

    # --- پردازش پیام‌ها و فایل‌ها ---
    gemini_messages = []
    for msg in data.get("messages", []):
        role = "model" if msg.get("role") == "assistant" else msg.get("role")
        processed_parts = []
        for part in msg.get("parts", []):
            if part.get("text"):
                processed_parts.append({"text": part["text"]})
            if part.get("base64Data") and part.get("mimeType"):
                mime_type = part["mimeType"]
                # هندل کردن فایل Word
                if mime_type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
                    try:
                        decoded = base64.b64decode(part["base64Data"])
                        doc = docx.Document(io.BytesIO(decoded))
                        text = "\n".join([p.text for p in doc.paragraphs])
                        processed_parts.append({"text": f"محتوای فایل کاربر:\n{text}"})
                    except:
                        processed_parts.append({"text": "(فایل Word قابل خواندن نبود)"})
                else:
                    processed_parts.append({"inline_data": {"mime_type": mime_type, "data": part["base64Data"]}})
        
        if processed_parts:
            if gemini_messages and gemini_messages[-1]["role"] == role:
                gemini_messages[-1]["parts"].extend(processed_parts)
            else:
                gemini_messages.append({"role": role, "parts": processed_parts})

    # اگر پیامی نبود تمام کن
    if not any(msg['role'] == 'user' for msg in gemini_messages):
        return Response("data: [DONE]\n\n", mimetype='text/event-stream')

    # --- تابع اصلی استریم با مدیریت خطای پیشرفته ---
    @stream_with_context
    def stream_response_generator():
        # تعداد دفعاتی که کلا تلاش میکنیم (مثلا 3 دور کامل روی همه کلیدها)
        # این باعث میشه کاربر احساس نکنه سرور قطع شده، فقط فکر میکنه داره فکر میکنه
        MAX_GLOBAL_RETRIES = len(GEMINI_API_KEYS) * 3 if len(GEMINI_API_KEYS) > 0 else 1
        
        # فلگ برای اینکه بفهمیم بالاخره موفق شدیم یا نه
        success = False

        for attempt in range(MAX_GLOBAL_RETRIES):
            try:
                api_key, idx = get_next_key_with_index()
                
                # تنظیمات تایم اوت:
                # 5 ثانیه برای اتصال (اگر کلید خراب بود سریع رد شو)
                # 100 ثانیه برای خواندن (اگر فایل سنگین بود صبر کن)
                response = requests.post(
                    f"https://generativelanguage.googleapis.com/v1beta/models/{GEMINI_MODEL_NAME}:streamGenerateContent?key={api_key}&alt=sse",
                    json={
                        "contents": gemini_messages,
                        "systemInstruction": {"parts": [{"text": system_instruction}]},
                        "generationConfig": {
                            "temperature": 0.7,
                            "thinking_config": {"include_thoughts": True} if show_thoughts else {}
                        }
                    },
                    stream=True,
                    timeout=(5, 100) 
                )

                # اگر کد 200 نبود، یعنی این کلید مشکل داره.
                # نکته مهم: اینجا هیچ چیزی به کاربر نمیفرستیم. `continue` میکنیم تا بره کلید بعدی.
                if response.status_code != 200:
                    response.close()
                    continue 

                # حالا چک میکنیم که آیا واقعا دیتایی میاد؟
                # این خط iterator رو میگیره اما هنوز دانلود نکرده
                line_iter = response.iter_lines()
                
                # سعی میکنیم اولین خط رو بگیریم.
                # اگر اینجا ارور بده (مثل 502 وسط کار)، میره توی except و کلید بعدی رو تست میکنه
                # پس کاربر هنوز چیزی ندیده.
                first_chunk_found = False
                
                for line in line_iter:
                    if line:
                        decoded_line = line.decode('utf-8')
                        if decoded_line.startswith('data: '):
                            try:
                                json_data = json.loads(decoded_line[6:])
                                # اگر جیسون ولید بود، یعنی اتصال درسته.
                                # حالا شروع میکنیم به فرستادن به کاربر
                                first_chunk_found = True
                                
                                # اینجا دیگه تسلیم میشیم و شروع میکنیم به ارسال به کلاینت
                                # چون مطمئن شدیم این کلید سالمه
                                success = True
                                
                                # پردازش معمولی
                                parts = json_data.get("candidates", [{}])[0].get("content", {}).get("parts", [])
                                for part in parts:
                                    if part.get("text"):
                                        is_thought = part.get("thought") is True
                                        if show_thoughts and is_thought:
                                            yield f"data: {json.dumps({'type': 'thought', 'content': part['text']})}\n\n"
                                        elif not is_thought:
                                            yield f"data: {json.dumps({'choices': [{'delta': {'content': part['text']}}]})}\n\n"
                                            
                            except:
                                pass # ایگنور کردن خطاهای جزئی جیسون

                # اگر حلقه تمام شد و ما دیتایی فرستادیم، یعنی کار تمامه
                if success:
                    return # خروج از کل تابع (پایان استریم)

            except Exception as e:
                # هر خطایی (قطع نت، تایم اوت، ارور گوگل)
                # فقط لاگ کن و برو کلید بعدی
                # کاربر چیزی حس نمیکنه، فقط اسپینر میچرخه
                logging.error(f"Retry {attempt}: {e}")
                time.sleep(0.2) # وقفه کوتاه
                continue

        # === اگر به اینجا رسیدیم یعنی تمام کلیدها تست شدند و هیچکدام کار نکردند ===
        # به جای اینکه ارور قرمز بفرستیم، یک متن معمولی میفرستیم
        # این باعث میشه کاربر فکر کنه ربات جواب داده ولی نتونسته
        if not success:
            fallback_msg = "🤔 سیستم کمی شلوغ است و نتوانستم پاسخ را کامل کنم. لطفاً دوباره تلاش کنید یا فایلتان را بررسی کنید."
            yield f"data: {json.dumps({'choices': [{'delta': {'content': fallback_msg}}]})}\n\n"

    return Response(stream_response_generator(), mimetype='text/event-stream')

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=os.environ.get("PORT", 7860))

# --- END OF FILE app.py ---