from typing import Optional, Union
from volcenginesdkarkruntime import AsyncArk
from openai import AsyncOpenAI

_ark_client: Optional[AsyncArk] = None
_aicc_client: Optional[AsyncOpenAI] = None
_is_aicc: bool = False

def set_ark_client(client: AsyncArk):
    global _ark_client
    _ark_client = client

def get_ark_client() -> Optional[AsyncArk]:
    return _ark_client

def set_aicc_client(client: AsyncOpenAI):
    global _aicc_client
    _aicc_client = client

def get_aicc_client() -> Optional[AsyncOpenAI]:
    return _aicc_client

def set_is_aicc(value: bool):
    global _is_aicc
    _is_aicc = value

def get_is_aicc() -> bool:
    return _is_aicc

def get_client() -> Union[AsyncArk, AsyncOpenAI, None]:
    """返回当前活跃的客户端（AICC 或 Ark）"""
    if _is_aicc:
        return _aicc_client
    return _ark_client

def cleanup_ark_client():
    global _ark_client
    if _ark_client:
        _ark_client = None

def cleanup_aicc_client():
    global _aicc_client
    if _aicc_client:
        _aicc_client = None

def cleanup_all():
    cleanup_ark_client()
    cleanup_aicc_client()