import os
import time
import logging
import requests
import concurrent.futures
from datetime import datetime, timedelta
from apscheduler.schedulers.background import BackgroundScheduler
from flask import Flask, request, jsonify, Response, stream_with_context

logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s - %(levelname)s - %(message)s')

API_ENDPOINT = "https://api.siliconflow.cn/v1/user/info"
TEST_MODEL_ENDPOINT = "https://api.siliconflow.cn/v1/chat/completions"
MODELS_ENDPOINT = "https://api.siliconflow.cn/v1/models"

app = Flask(__name__)

all_models = []
free_models = []

invalid_keys_global = []
free_keys_global = []
unverified_keys_global = []
valid_keys_global = []

# 使用 ThreadPoolExecutor 管理并发
executor = concurrent.futures.ThreadPoolExecutor(max_workers=20)

# 为每个模型维护一个独立的 key 索引
model_key_indices = {}

def get_credit_summary(api_key):
    """
    使用 API 密钥获取额度信息。
    """
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    try:
        response = requests.get(API_ENDPOINT, headers=headers)
        response.raise_for_status()
        data = response.json().get("data", {})
        total_balance = data.get("totalBalance", 0)
        return {"total_balance": float(total_balance)}
    except requests.exceptions.RequestException as e:
        logging.error(f"获取额度信息失败，API Key：{api_key}，错误信息：{e}")
        return None
    except (KeyError, TypeError) as e:
        logging.error(f"解析额度信息失败，API Key：{api_key}，错误信息：{e}")
        return None
    except ValueError as e:
        logging.error(f"total_balance 无法转换为浮点数，API Key：{api_key}，错误信息：{e}")
        return None

FREE_MODEL_TEST_KEY = "sk-bmjbjzleaqfgtqfzmcnsbagxrlohriadnxqrzfocbizaxukw"

def test_model_availability(api_key, model_name):
    """
    测试指定的模型是否可用。
    """
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    try:
        response = requests.post(TEST_MODEL_ENDPOINT,
                                 headers=headers,
                                 json={
                                     "model": model_name,
                                     "messages": [{"role": "user", "content": "hi"}],
                                     "max_tokens": 10,
                                     "stream": False
                                 },
                                 timeout=10)
        if response.status_code == 429 or response.status_code == 200 :
            return True
        else:
            return False
    except requests.exceptions.RequestException as e:
        logging.error(f"测试模型 {model_name} 可用性失败，API Key：{api_key}，错误信息：{e}")
        return False

def refresh_models():
    """
    刷新模型列表和免费模型列表。
    """
    global all_models, free_models

    all_models = get_all_models(FREE_MODEL_TEST_KEY)
    free_models = []

    with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
        future_to_model = {executor.submit(test_model_availability, FREE_MODEL_TEST_KEY, model): model for model in all_models}
        for future in concurrent.futures.as_completed(future_to_model):
            model = future_to_model[future]
            try:
                is_free = future.result()
                if is_free:
                    free_models.append(model)
            except Exception as exc:
                logging.error(f"模型 {model} 测试生成异常: {exc}")

    logging.info(f"所有模型列表：{all_models}")
    logging.info(f"免费模型列表：{free_models}")

def load_keys():
    """
    从环境变量中加载 keys，并根据额度和模型可用性进行分类，然后记录到日志中。
    使用线程池并发处理每个 key。
    """
    keys_str = os.environ.get("KEYS")
    test_model = os.environ.get("TEST_MODEL", "Pro/google/gemma-2-9b-it")

    if keys_str:
        keys = [key.strip() for key in keys_str.split(',')]
        logging.info(f"加载的 keys：{keys}")

        with concurrent.futures.ThreadPoolExecutor(max_workers=20) as executor:
            future_to_key = {executor.submit(process_key, key, test_model): key for key in keys}

            invalid_keys = []
            free_keys = []
            unverified_keys = []
            valid_keys = []

            for future in concurrent.futures.as_completed(future_to_key):
                key = future_to_key[future]
                try:
                    key_type = future.result()
                    if key_type == "invalid":
                        invalid_keys.append(key)
                    elif key_type == "free":
                        free_keys.append(key)
                    elif key_type == "unverified":
                        unverified_keys.append(key)
                    elif key_type == "valid":
                        valid_keys.append(key)
                except Exception as exc:
                    logging.error(f"处理 KEY {key} 生成异常: {exc}")

        logging.info(f"无效 KEY：{invalid_keys}")
        logging.info(f"免费 KEY：{free_keys}")
        logging.info(f"未实名 KEY：{unverified_keys}")
        logging.info(f"有效 KEY：{valid_keys}")

        global invalid_keys_global, free_keys_global, unverified_keys_global, valid_keys_global
        invalid_keys_global = invalid_keys
        free_keys_global = free_keys
        unverified_keys_global = unverified_keys
        valid_keys_global = valid_keys

    else:
        logging.warning("环境变量 KEYS 未设置。")

def process_key(key, test_model):
    """
    处理单个 key，判断其类型。
    """
    credit_summary = get_credit_summary(key)
    if credit_summary is None:
        return "invalid"
    else:
        total_balance = credit_summary.get("total_balance", 0)
        if total_balance <= 0:
            return "free"
        else:
            if test_model_availability(key, test_model):
                return "valid"
            else:
                return "unverified"

def get_all_models(api_key):
    """
    获取所有模型列表。
    """
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    try:
        response = requests.get(MODELS_ENDPOINT, headers=headers, params={"sub_type": "chat"})
        response.raise_for_status()
        data = response.json()
        if isinstance(data, dict) and 'data' in data and isinstance(data['data'], list):
            return [model.get("id") for model in data["data"] if isinstance(model, dict) and "id" in model]
        else:
            logging.error("获取模型列表失败：响应数据格式不正确")
            return []
    except requests.exceptions.RequestException as e:
        logging.error(f"获取模型列表失败，API Key：{api_key}，错误信息：{e}")
        return []
    except (KeyError, TypeError) as e:
        logging.error(f"解析模型列表失败，API Key：{api_key}，错误信息：{e}")
        return []

def determine_request_type(model_name):
    """
    根据用户请求的模型判断请求类型。
    """
    if model_name in free_models:
        return "free"
    elif model_name in all_models:
        return "paid"
    else:
        return "unknown"

def select_key(request_type, model_name):
    """
    根据请求类型和模型名称选择合适的 KEY，并实现轮询和重试机制。
    """
    if request_type == "free":
        available_keys = free_keys_global + unverified_keys_global + valid_keys_global
    elif request_type == "paid":
        available_keys = unverified_keys_global + valid_keys_global
    else:
        available_keys = free_keys_global + unverified_keys_global + valid_keys_global

    if not available_keys:
        return None

    # 获取当前模型的索引，如果没有则初始化为 0
    current_index = model_key_indices.get(model_name, 0)

    # 轮询并重试
    for _ in range(len(available_keys)):
        key = available_keys[current_index % len(available_keys)]
        current_index += 1
        
        # 检查 KEY 是否有效
        if key_is_valid(key, request_type):
            # 更新模型索引并返回 KEY
            model_key_indices[model_name] = current_index
            return key
        else:
            logging.warning(f"KEY {key} 无效或达到限制，尝试下一个 KEY")

    # 所有 KEY 都尝试过，重置索引并返回 None
    model_key_indices[model_name] = 0
    return None

def key_is_valid(key, request_type):
    """
    检查 KEY 是否有效，根据不同的请求类型进行不同的检查。
    """
    if request_type == "invalid":
        return False  # 无效 KEY 始终返回 False
    
    credit_summary = get_credit_summary(key)
    if credit_summary is None:
        return False
    
    total_balance = credit_summary.get("total_balance", 0)
    
    if request_type == "free":
        return True # 免费 KEY，只要能获取到信息，就认为是有效的
    elif request_type == "paid" or request_type == "unverified":
        return total_balance > 0 # 付费 KEY 或未实名 KEY，需要余额大于 0 才是有效的
    else:
        return False # 未知类型，返回 False

def check_authorization(request):
    """
    检查请求头中的 Authorization 字段是否匹配环境变量 AUTHORIZATION_KEY。
    """
    authorization_key = os.environ.get("AUTHORIZATION_KEY")
    if not authorization_key:
        logging.warning("环境变量 AUTHORIZATION_KEY 未设置，请设置后重试。")
        return False

    auth_header = request.headers.get('Authorization')
    if not auth_header:
        logging.warning("请求头中缺少 Authorization 字段。")
        return False

    if auth_header != f"Bearer {authorization_key}":
        logging.warning(f"无效的 Authorization 密钥：{auth_header}")
        return False

    return True

scheduler = BackgroundScheduler()
scheduler.add_job(load_keys, 'interval', hours=1)
scheduler.add_job(refresh_models, 'interval', minutes=10)

@app.route('/')
def index():
    return "<h1>Welcome to SiliconFlow</h1>"

@app.route('/check_tokens', methods=['POST'])
def check_tokens():
    """
    处理前端发送的 Token 检测请求。
    使用线程池并发处理每个 token。
    """
    tokens = request.json.get('tokens', [])
    test_model = os.environ.get("TEST_MODEL", "Pro/google/gemma-2-9b-it")

    with concurrent.futures.ThreadPoolExecutor(max_workers=20) as executor:
        future_to_token = {executor.submit(process_key, token, test_model): token for token in tokens}

        results = []
        for future in concurrent.futures.as_completed(future_to_token):
            token = future_to_token[future]
            try:
                key_type = future.result()
                credit_summary = get_credit_summary(token)
                balance = credit_summary.get("total_balance", 0) if credit_summary else 0
                if key_type == "invalid":
                    results.append({"token": token, "type": "无效 KEY", "balance": balance, "message": "无法获取额度信息"})
                elif key_type == "free":
                    results.append({"token": token, "type": "免费 KEY", "balance": balance, "message": "额度不足"})
                elif key_type == "unverified":
                    results.append({"token": token, "type": "未实名 KEY", "balance": balance, "message": "无法使用指定模型"})
                elif key_type == "valid":
                    results.append({"token": token, "type": "有效 KEY", "balance": balance, "message": "可以使用指定模型"})
            except Exception as exc:
                logging.error(f"处理 Token {token} 生成异常: {exc}")

    return jsonify(results)

@app.route('/handsome/v1/chat/completions', methods=['POST'])
def handsome_chat_completions():
    """
    处理 /handsome/v1/chat/completions 路由的请求，添加鉴权，并实现 KEY 的轮询和重试机制。
    """
    if not check_authorization(request):
        return jsonify({"error": "Unauthorized"}), 401

    data = request.get_json()
    if not data or 'model' not in data:
        return jsonify({"error": "Invalid request data"}), 400

    model_name = data['model']
    request_type = determine_request_type(model_name)
    api_key = select_key(request_type, model_name)

    if not api_key:
        return jsonify({"error": "No available API key for this request type or all keys have reached their limits"}), 429

    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }

    try:
        response = requests.post(
            TEST_MODEL_ENDPOINT,
            headers=headers,
            json=data,
            stream=data.get("stream", False),
            timeout=60
        )

        if response.status_code == 429:
            return jsonify(response.json()), 429

        if data.get("stream", False):
            return Response(stream_with_context(response.iter_content(chunk_size=1024)), content_type=response.headers['Content-Type'])
        else:
            response.raise_for_status()
            return jsonify(response.json())
    except requests.exceptions.RequestException as e:
        return jsonify({"error": str(e)}), 500

@app.route('/handsome/v1/models', methods=['GET'])
def list_models():
    if not check_authorization(request):
        return jsonify({"error": "Unauthorized"}), 401

    return jsonify({
        "data": [{"id": model, "object": "model"} for model in all_models],
        "free_models": free_models
    })

def get_billing_info():
    """
    获取所有KEY的额度信息。
    使用线程池并发处理每个 key。
    """
    keys = valid_keys_global + unverified_keys_global
    total_balance = 0

    with concurrent.futures.ThreadPoolExecutor(max_workers=20) as executor:
        futures = [executor.submit(get_credit_summary, key) for key in keys]

        for future in concurrent.futures.as_completed(futures):
            try:
                credit_summary = future.result()
                if credit_summary:
                    total_balance += credit_summary.get("total_balance", 0)
            except Exception as exc:
                logging.error(f"获取额度信息生成异常: {exc}")

    return total_balance

@app.route('/handsome/v1/dashboard/billing/usage', methods=['GET'])
def billing_usage():
    if not check_authorization(request):
        return jsonify({"error": "Unauthorized"}), 401

    end_date = datetime.now()
    start_date = end_date - timedelta(days=30)

    daily_usage = []
    current_date = start_date
    while current_date <= end_date:
        daily_usage.append({
            "timestamp": int(current_date.timestamp()),
            "daily_usage": 0
        })
        current_date += timedelta(days=1)

    return jsonify({
        "object": "list",
        "data": daily_usage,
        "total_usage": 0
    })

@app.route('/handsome/v1/dashboard/billing/subscription', methods=['GET'])
def billing_subscription():
    if not check_authorization(request):
        return jsonify({"error": "Unauthorized"}), 401

    total_balance = get_billing_info()

    return jsonify({
        "object": "billing_subscription",
        "has_payment_method": False,
        "canceled": False,
        "canceled_at": None,
        "delinquent": None,
        "access_until": int(datetime(9999, 12, 31).timestamp()),
        "soft_limit": 0,
        "hard_limit": total_balance,
        "system_hard_limit": total_balance,
        "soft_limit_usd": 0,
        "hard_limit_usd": total_balance,
        "system_hard_limit_usd": total_balance,
        "plan": {
            "name": "SiliconFlow API",
            "id": "siliconflow-api"
        },
        "account_name": "SiliconFlow User",
        "po_number": None,
        "billing_email": None,
        "tax_ids": [],
        "billing_address": None,
        "business_address": None
    })

if __name__ == '__main__':
    logging.info(f"环境变量：{os.environ}")

    invalid_keys_global = []
    free_keys_global = []
    unverified_keys_global = []
    valid_keys_global = []

    scheduler.start()

    load_keys()
    logging.info("首次加载 keys 已手动触发执行")

    refresh_models()
    logging.info("首次刷新模型列表已手动触发执行")

    app.run(debug=False, host='0.0.0.0', port=int(os.environ.get('PORT', 7860)))