"""
模型对话 API 路由
"""
import os
import pymysql
import yaml
import requests
import concurrent.futures
from flask import Blueprint, request, jsonify

# 获取项目根目录
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 创建蓝图
model_chat_bp = Blueprint('model_chat', __name__, url_prefix='/api/model-chat')


def get_db_connection():
    """获取数据库连接"""
    CONFIG_PATH = os.path.join(PROJECT_ROOT, 'config.yaml')
    with open(CONFIG_PATH, 'r', encoding='utf-8') as f:
        CONFIG = yaml.safe_load(f)
    db_config = CONFIG['database']
    return pymysql.connect(
        host=db_config['host'],
        port=db_config['port'],
        user=db_config['username'],
        password=db_config['password'],
        database=db_config['name'],
        charset=db_config.get('charset', 'utf8mb4'),
        cursorclass=pymysql.cursors.DictCursor
    )


def generic_get_by_id(table_name, id_val):
    """按ID查询"""
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute(f"SELECT * FROM {table_name} WHERE id = %s", (id_val,))
    result = cursor.fetchone()
    cursor.close()
    conn.close()
    return result


def call_api_model(model_config, messages, temperature, max_tokens):
    """调用API模型（OpenAI兼容格式）"""
    api_url = model_config.get('api_url')
    api_key = model_config.get('api_key')
    model_name = model_config.get('model_name', '')

    # 构造OpenAI兼容的完整URL
    # 支持: https://api.openai.com/v1/chat/completions 或 https://api.example.com/v1
    # 如果URL已经包含 /chat/completions 则直接使用，否则追加
    if '/chat/completions' in api_url:
        full_url = api_url
    else:
        # 去掉末尾的斜杠，然后追加 /chat/completions
        full_url = api_url.rstrip('/') + '/chat/completions'

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

    payload = {
        'model': model_name,
        'messages': messages,
        'temperature': temperature,
        'max_tokens': max_tokens
    }

    try:
        response = requests.post(full_url, headers=headers, json=payload, timeout=120)
        response.raise_for_status()
        result = response.json()

        if 'choices' in result and len(result['choices']) > 0:
            return {
                'success': True,
                'content': result['choices'][0]['message'].get('content', '')
            }
        return {'success': False, 'error': 'API返回格式异常'}
    except requests.exceptions.RequestException as e:
        return {'success': False, 'error': str(e)}


def call_local_model(model_config, messages, temperature, max_tokens):
    """调用本地模型（通过vLLM OpenAI兼容API）"""
    api_url = model_config.get('path')  # 本地模型path字段存储API地址
    model_name = model_config.get('model_name', '')

    if not api_url:
        return {'success': False, 'error': '本地模型API地址未配置'}

    headers = {'Content-Type': 'application/json'}

    payload = {
        'model': model_name,
        'messages': messages,
        'temperature': temperature,
        'max_tokens': max_tokens
    }

    try:
        response = requests.post(api_url, headers=headers, json=payload, timeout=120)
        response.raise_for_status()
        result = response.json()

        if 'choices' in result and len(result['choices']) > 0:
            return {
                'success': True,
                'content': result['choices'][0]['message'].get('content', '')
            }
        return {'success': False, 'error': 'API返回格式异常'}
    except requests.exceptions.RequestException as e:
        return {'success': False, 'error': str(e)}


@model_chat_bp.route('', methods=['POST'])
def model_chat():
    """模型对话接口"""
    data = request.json
    model_id = data.get('model_id')
    system_prompt = data.get('system_prompt', '')
    user_question = data.get('user_question')
    temperature = data.get('temperature', 0.7)
    max_tokens = data.get('max_tokens', 2048)

    if not model_id:
        return jsonify({'code': 1, 'message': '缺少模型ID'})
    if not user_question:
        return jsonify({'code': 1, 'message': '缺少用户提问'})

    # 获取模型配置
    model = generic_get_by_id('model_manage', model_id)
    if not model:
        return jsonify({'code': 1, 'message': '模型不存在'})

    # 构建消息
    messages = []
    if system_prompt:
        messages.append({'role': 'system', 'content': system_prompt})
    messages.append({'role': 'user', 'content': user_question})

    # 根据模型类型调用
    if model.get('model_source') == 'api':
        result = call_api_model(model, messages, temperature, max_tokens)
    else:
        result = call_local_model(model, messages, temperature, max_tokens)

    if result.get('success'):
        return jsonify({
            'code': 0,
            'data': {
                'model_id': model_id,
                'model_name': model.get('name'),
                'response': result['content']
            }
        })
    else:
        return jsonify({'code': 1, 'message': result.get('error', '调用失败')})


@model_chat_bp.route('/batch', methods=['POST'])
def model_chat_batch():
    """批量模型对话接口（并发调用多个模型）"""
    data = request.json
    model_ids = data.get('model_ids', [])
    system_prompt = data.get('system_prompt', '')
    user_question = data.get('user_question')
    temperature = data.get('temperature', 0.7)
    max_tokens = data.get('max_tokens', 2048)

    if not model_ids:
        return jsonify({'code': 1, 'message': '缺少模型ID列表'})
    if not user_question:
        return jsonify({'code': 1, 'message': '缺少用户提问'})

    def call_single_model(model_id):
        model = generic_get_by_id('model_manage', model_id)
        if not model:
            return {'model_id': model_id, 'success': False, 'error': '模型不存在'}

        messages = []
        if system_prompt:
            messages.append({'role': 'system', 'content': system_prompt})
        messages.append({'role': 'user', 'content': user_question})

        if model.get('model_source') == 'api':
            result = call_api_model(model, messages, temperature, max_tokens)
        else:
            result = call_local_model(model, messages, temperature, max_tokens)

        return {
            'model_id': model_id,
            'model_name': model.get('name'),
            'success': result.get('success', False),
            'response': result.get('content', ''),
            'error': result.get('error', '')
        }

    # 并发调用所有模型
    results = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=min(len(model_ids), 4)) as executor:
        future_to_model = {executor.submit(call_single_model, mid): mid for mid in model_ids}
        for future in concurrent.futures.as_completed(future_to_model):
            results.append(future.result())

    return jsonify({'code': 0, 'data': results})
