# --- 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
import requests
import docx

# ================== بخش تنظیمات لاگ‌نویسی (بدون تغییر) ==================

class NoGrpcFilter(logging.Filter):
    def filter(self, record):
        return not record.getMessage().startswith('ALTS creds ignored.')

def setup_logging():
    log_format = '[%(asctime)s] [%(levelname)s]: %(message)s'
    date_format = '%Y-%m-%d %H:%M:%S'
    formatter = logging.Formatter(log_format, datefmt=date_format)
    
    root_logger = logging.getLogger()
    if root_logger.hasHandlers():
        root_logger.handlers.clear()
    
    console_handler = logging.StreamHandler()
    console_handler.setFormatter(formatter)
    console_handler.addFilter(NoGrpcFilter())

    root_logger.addHandler(console_handler)
    root_logger.setLevel(logging.INFO)

setup_logging()
app = Flask(__name__)

# ================== بخش پیکربندی Gemini (با تغییر) ==================

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()]

if not GEMINI_API_KEYS:
    logging.critical("هشدار: هیچ کلید API برای Gemini در Secrets تنظیم نشده است! (ALL_GEMINI_API_KEYS)")

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:
            raise ValueError("لیست کلیدهای API خالی است.")
        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:
        # اگر هیچ کلیدی نباشد چاره‌ای جز خطا نیست، اما این حالت نادری است
        return Response("data: [DONE]\n\n", mimetype='text/event-stream')

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

    # بخش پردازش پیام‌ها و فایل DOCX
    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"]
                
                if mime_type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
                    try:
                        decoded_data = base64.b64decode(part["base64Data"])
                        file_stream = io.BytesIO(decoded_data)
                        document = docx.Document(file_stream)
                        full_text = "\n".join([para.text for para in document.paragraphs])
                        
                        final_text_part = f"کاربر یک فایل Word آپلود کرد. محتوای متنی آن به شرح زیر است:\n\n---\n\n{full_text}\n\n---"
                        processed_parts.append({"text": final_text_part})
                        logging.info("فایل DOCX با موفقیت پردازش و متن آن استخراج شد.")

                    except Exception as e:
                        logging.error(f"خطا در پردازش فایل DOCX: {e}")
                        processed_parts.append({"text": "[خطا: امکان پردازش فایل Word وجود نداشت.]"})
                
                else:
                    processed_parts.append({"inline_data": {"mime_type": part["mimeType"], "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')
        
    def stream_response():
        # تلاش برای تعداد زیادی بار (عملا تا وقتی یک کلید سالم پیدا شود)
        # ضرب در 2 یعنی هر کلید را دو بار شانس می دهیم اگر شبکه قطع و وصل شد
        max_attempts = len(GEMINI_API_KEYS) * 3 
        
        for attempt in range(max_attempts):
            try:
                api_key, key_index = get_next_key_with_index()
                # logging.info(f"تلاش با کلید شماره {key_index + 1}...")

                api_endpoint = f"https://generativelanguage.googleapis.com/v1beta/models/{GEMINI_MODEL_NAME}:streamGenerateContent?key={api_key}&alt=sse"
                
                payload = {
                    "contents": gemini_messages,
                    "systemInstruction": {"parts": [{"text": system_instruction}]},
                    "tools": [{"google_search": {}}],
                    "generationConfig": {
                        "temperature": 0.7,
                    }
                }
                
                if show_thoughts:
                    payload["generationConfig"]["thinking_config"] = {
                        "include_thoughts": True
                    }

                # *** نکته مهم: Timeout بسیار کوتاه برای اتصال، تا سریع برود کلید بعدی ***
                # connect=4: اگر 4 ثانیه وصل نشد، ولش کن برو بعدی
                # read=20: اگر وسط کار 20 ثانیه دیتا نیامد، ولش کن برو بعدی
                with requests.post(api_endpoint, json=payload, stream=True, timeout=(4, 20)) as response:
                    
                    # اگر ارور 429 (محدودیت) یا هر ارور دیگری داد، سریع برو بعدی
                    if response.status_code != 200:
                        logging.warning(f"کلید {key_index + 1} با وضعیت {response.status_code} پاسخ نداد. تلاش با کلید بعدی...")
                        continue 
                    
                    # اگر موفق شد، شروع به ارسال کن و از تابع خارج شو (return)
                    for line in response.iter_lines():
                        if line:
                            decoded_line = line.decode('utf-8')
                            if decoded_line.startswith('data: '):
                                try:
                                    chunk_data = json.loads(decoded_line[6:])
                                    parts = chunk_data.get("candidates", [{}])[0].get("content", {}).get("parts", [])
                                    
                                    for part in parts:
                                        if "text" not in part or not part["text"]:
                                            continue
                                        is_a_thought = part.get("thought") is True
                                        if show_thoughts and is_a_thought:
                                            thought_payload = {"type": "thought", "content": part["text"]}
                                            yield f"data: {json.dumps(thought_payload)}\n\n"
                                        elif not is_a_thought:
                                            sse_payload = {"choices": [{"delta": {"content": part["text"]}}]}
                                            yield f"data: {json.dumps(sse_payload)}\n\n"
                                except (json.JSONDecodeError, IndexError, KeyError):
                                    continue
                    
                    # اگر استریم بدون خطا تمام شد، کار تمام است
                    return

            except Exception as e:
                # هر نوع خطایی (تایم اوت، شبکه و ...) رخ داد، فقط لاگ کن و برو کلید بعدی
                # هیچ خطایی به کاربر نشان نده
                logging.warning(f"خطا در کلید {key_index + 1}: {e} - رفتن به کلید بعدی")
                continue 
        
        # اگر همه کلیدها تست شدند و نشد (خیلی بعید)، یک پیام خالی بفرست که ارور نده
        logging.critical("هیچ کلیدی کار نکرد.")
        # به جای ارور، یک پیام ساده میفرستیم که برنامه کرش نکند
        end_payload = {"choices": [{"delta": {"content": "..."}}]} 
        yield f"data: {json.dumps(end_payload)}\n\n"

    return Response(stream_response(), 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 ---