# app/core/gemini_client_manager.py

import os  # 导入 os 模块用于读取环境变量
import itertools  # 导入 itertools 模块用于循环 cookie
import logging
from typing import Dict, Any, List, Optional, Tuple
from gemini_webapi import GeminiClient
import gemini_webapi.exceptions # 导入 gemini_webapi 的异常模块
import asyncio # 导入 asyncio 模块用于异步操作
from app.core.session_manager import chat_sessions # 导入 chat_sessions

logger = logging.getLogger(__name__)

# 新的环境变量读取方式
GEMINI_PSID_COOKIES_STR = os.getenv("GEMINI_PSID_COOKIES", "")
GEMINI_PSIDTS_COOKIES_STR = os.getenv("GEMINI_PSIDTS_COOKIES", "")

raw_psid_cookies = [c.strip() for c in GEMINI_PSID_COOKIES_STR.split(",") if c.strip()]
raw_psidts_cookies = [
    c.strip() for c in GEMINI_PSIDTS_COOKIES_STR.split(",") if c.strip()
]

gemini_cookie_pairs: List[Tuple[str, str]] = []
if raw_psid_cookies and raw_psidts_cookies:
    if len(raw_psid_cookies) == len(raw_psidts_cookies):
        valid_pairs = []
        for sid, sidts in zip(raw_psid_cookies, raw_psidts_cookies):
            if sid and sidts:  # Both must be non-empty
                valid_pairs.append((sid, sidts))
            else:
                logger.warning(
                    f"检测到无效的 Cookie 对 (一个或两个值为空): SID='{sid}', SIDTS='{sidts}'. 此对将被忽略。"
                )

        gemini_cookie_pairs = valid_pairs
        if gemini_cookie_pairs:
            logger.info(
                f"成功加载并验证 {len(gemini_cookie_pairs)} 组 Gemini Cookie 对。"
            )
        else:
            logger.warning(
                "所有提供的 Cookie 对均无效 (一个或两个值为空) 或 GEMINI_PSID_COOKIES / GEMINI_PSIDTS_COOKIES 列表长度匹配但内容无效。"
            )
    else:
        logger.warning(
            "GEMINI_PSID_COOKIES 和 GEMINI_PSIDTS_COOKIES 环境变量中的 Cookie 数量不匹配。请检查配置。"
        )
else:
    if not raw_psid_cookies and not raw_psidts_cookies:
        logger.warning(
            "GEMINI_PSID_COOKIES 和 GEMINI_PSIDTS_COOKIES 环境变量均未设置或为空。Gemini API 调用可能会失败。"
        )
    elif not raw_psid_cookies:
        logger.warning(
            "GEMINI_PSID_COOKIES 环境变量未设置或为空，但 GEMINI_PSIDTS_COOKIES 可能已设置。Cookie 对不完整或无效。"
        )
    else:  # not raw_psidts_cookies
        logger.warning(
            "GEMINI_PSIDTS_COOKIES 环境变量未设置或为空，但 GEMINI_PSID_COOKIES 可能已设置。Cookie 对不完整或无效。"
        )

# 全局变量，用于存储 cookie 轮询器和客户端实例
cookie_cycler = itertools.cycle(gemini_cookie_pairs) if gemini_cookie_pairs else None
gemini_clients: Dict[str, GeminiClient] = {} # 存储 {auth_token: GeminiClient实例}

def reload_gemini_cookies():
    """
    重新加载 Gemini Cookies 环境变量并重新初始化 cookie 轮询器。
    此函数应在需要刷新 cookies 时调用。
    """
    global gemini_cookie_pairs, cookie_cycler, gemini_clients # 声明为全局变量以便修改

    logger.info("开始重新加载 Gemini Cookies...")

    # 重新读取环境变量
    new_psid_cookies_str = os.getenv("GEMINI_PSID_COOKIES", "")
    new_psidts_cookies_str = os.getenv("GEMINI_PSIDTS_COOKIES", "")

    new_raw_psid_cookies = [c.strip() for c in new_psid_cookies_str.split(",") if c.strip()]
    new_raw_psidts_cookies = [c.strip() for c in new_psidts_cookies_str.split(",") if c.strip()]

    new_gemini_cookie_pairs = []
    if new_raw_psid_cookies and new_raw_psidts_cookies:
        if len(new_raw_psid_cookies) == len(new_raw_psidts_cookies):
            valid_pairs = []
            for sid, sidts in zip(new_raw_psid_cookies, new_raw_psidts_cookies):
                if sid and sidts:
                    valid_pairs.append((sid, sidts))
                else:
                    logger.warning(
                        f"重新加载时检测到无效的 Cookie 对 (一个或两个值为空): SID='{sid}', SIDTS='{sidts}'. 此对将被忽略。"
                    )
            new_gemini_cookie_pairs = valid_pairs
            if new_gemini_cookie_pairs:
                logger.info(
                    f"成功重新加载并验证 {len(new_gemini_cookie_pairs)} 组 Gemini Cookie 对。"
                )
            else:
                logger.warning(
                    "重新加载时所有提供的 Cookie 对均无效或列表长度匹配但内容无效。"
                )
        else:
            logger.warning(
                "重新加载时 GEMINI_PSID_COOKIES 和 GEMINI_PSIDTS_COOKIES 环境变量中的 Cookie 数量不匹配。请检查配置。"
            )
    else:
        logger.warning(
            "重新加载时 GEMINI_PSID_COOKIES 和 GEMINI_PSIDTS_COOKIES 环境变量均未设置或为空。Gemini API 调用可能会失败。"
        )

    gemini_cookie_pairs = new_gemini_cookie_pairs
    cookie_cycler = itertools.cycle(gemini_cookie_pairs) if gemini_cookie_pairs else None

    # 清空所有现有的 GeminiClient 实例，强制下次请求时重新初始化
    gemini_clients.clear()
    logger.info("所有现有 GeminiClient 实例已清空，将在下次请求时重新初始化。")

    # 遍历所有会话，清空其 _chat_instance，强制重新初始化聊天实例
    for auth_token in chat_sessions:
        for session_id, session in chat_sessions[auth_token].items():
            if session._chat_instance:
                session._chat_instance = None
                logger.info(f"会话 {session_id} 的聊天实例已清空。")

    logger.info("Gemini Cookies 重新加载完成。")

def get_next_gemini_cookie_tuple() -> Optional[Tuple[str, str]]:
    """获取下一组可用的 Gemini Cookie (SID, SIDTS) 元组。"""
    if not cookie_cycler:
        logger.error(
            "Gemini Cookie 轮询器未初始化或为空。请检查 GEMINI_PSID_COOKIES 和 GEMINI_PSIDTS_COOKIES 环境变量。"
        )
        return None

    try:
        return next(cookie_cycler)
    except StopIteration:
        logger.error(
            "Gemini Cookie 轮询器意外耗尽。这可能表明初始化的 gemini_cookie_pairs 为空或所有元素均无效。"
        )
        return None
    except Exception as e:
        logger.error(f"获取下一组 Gemini Cookie 时发生意外错误: {e}", exc_info=True)
        return None

async def get_gemini_client(auth_token: str) -> GeminiClient:
    """
    获取或初始化与给定认证令牌关联的 GeminiClient 实例。
    如果客户端不存在，则创建一个新的。
    """
    client = gemini_clients.get(auth_token)

    if client:
        logger.info(f"认证令牌 {auth_token} 复用现有 GeminiClient 实例。")
        return client
    else:
        logger.info(f"认证令牌 {auth_token} 创建新的 GeminiClient 实例。")
        cookie_tuple = get_next_gemini_cookie_tuple()
        if not cookie_tuple:
            logger.error("未配置有效的 Gemini Cookies。请检查环境变量。")
            raise Exception("服务端未配置有效的 Gemini Cookies，无法连接到 Gemini API。")
        sid, sidts = cookie_tuple

        try:
            client = GeminiClient(sid, sidts)
            await client.init(timeout=30)  # 设置初始化超时时间
            gemini_clients[auth_token] = client  # 存储新的客户端实例
            logger.info(f"认证令牌 {auth_token} 成功初始化并存储 GeminiClient。")
            return client
        except Exception as client_init_error:
            logger.error(f"认证令牌 {auth_token} 初始化 GeminiClient 失败。")
            raise Exception("初始化 Gemini 客户端失败。请检查 GEMINI_PSID_COOKIES 和 GEMINI_PSIDTS_COOKIES 是否有效。")

# 在应用启动时调用一次，确保初始加载
reload_gemini_cookies()