"""
数据集管理 API 路由
"""
import io
import os
import time
import zipfile
from flask import Blueprint, request, jsonify, send_from_directory, Response
from werkzeug.utils import secure_filename

# 获取项目根目录
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
DATASET_FOLDER = os.path.join(PROJECT_ROOT, 'datasets')
ALLOWED_EXTENSIONS = {'jsonl', 'json', 'xls', 'xlsx'}

# 创建蓝图
datasets_bp = Blueprint('datasets', __name__, url_prefix='/api/dataset-manage')


def get_db_connection():
    """获取数据库连接"""
    import pymysql
    import yaml
    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 format_file_size(size_bytes):
    """格式化文件大小"""
    if size_bytes < 1024:
        return f"{size_bytes} B"
    elif size_bytes < 1024 * 1024:
        return f"{size_bytes / 1024:.1f} KB"
    elif size_bytes < 1024 * 1024 * 1024:
        return f"{size_bytes / (1024 * 1024):.1f} MB"
    else:
        return f"{size_bytes / (1024 * 1024 * 1024):.1f} GB"


def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


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


# ============ 数据集管理 CRUD ============

@datasets_bp.route('/<int:id>', methods=['GET'])
def get_dataset(id):
    """获取单个数据集详情"""
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute("SELECT * FROM dataset_manage WHERE id = %s", (id,))
    dataset = cursor.fetchone()

    if not dataset:
        cursor.close()
        conn.close()
        return jsonify({'code': 1, 'message': '数据集不存在'})

    # 获取关联的文件列表
    cursor.execute(
        "SELECT id, file_name, file_path, file_size, file_type, create_time FROM dataset_files WHERE dataset_id = %s ORDER BY create_time DESC",
        (id,)
    )
    files = cursor.fetchall()

    # 格式化文件大小
    for f in files:
        f['file_size_formatted'] = format_file_size(f['file_size'])

    dataset['files'] = files
    cursor.close()
    conn.close()

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


@datasets_bp.route('', methods=['GET'])
def get_datasets():
    """获取所有数据集"""
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute("SELECT * FROM dataset_manage ORDER BY create_time DESC")
    result = cursor.fetchall()
    cursor.close()
    conn.close()
    return jsonify({'code': 0, 'data': result})


@datasets_bp.route('', methods=['POST'])
def create_dataset():
    """创建数据集"""
    data = request.json
    try:
        conn = get_db_connection()
        cursor = conn.cursor()
        columns = ', '.join(data.keys())
        placeholders = ', '.join(['%s'] * len(data))
        sql = f"INSERT INTO dataset_manage ({columns}) VALUES ({placeholders})"
        cursor.execute(sql, list(data.values()))
        conn.commit()
        new_id = cursor.lastrowid
        cursor.close()
        conn.close()
        return jsonify({'code': 0, 'message': '创建成功', 'id': new_id})
    except Exception as e:
        return jsonify({'code': 1, 'message': f'创建失败: {str(e)}'})


@datasets_bp.route('/<int:id>', methods=['PUT'])
def update_dataset(id):
    """更新数据集"""
    data = request.json
    conn = get_db_connection()
    cursor = conn.cursor()
    set_clause = ', '.join([f"{k} = %s" for k in data.keys()])
    sql = f"UPDATE dataset_manage SET {set_clause} WHERE id = %s"
    values = list(data.values()) + [id]
    cursor.execute(sql, values)
    conn.commit()
    cursor.close()
    conn.close()
    return jsonify({'code': 0, 'message': '更新成功'})


@datasets_bp.route('/<int:id>', methods=['DELETE'])
def delete_dataset(id):
    """删除数据集"""
    conn = get_db_connection()
    cursor = conn.cursor()
    # 获取文件路径列表
    cursor.execute("SELECT file_path FROM dataset_files WHERE dataset_id = %s", (id,))
    files = cursor.fetchall()
    # 删除文件
    for f in files:
        file_path = f.get('file_path')
        if file_path and os.path.exists(file_path):
            try:
                os.remove(file_path)
            except Exception as e:
                print(f"删除文件失败: {file_path}, {e}")
    # 删除数据库记录
    cursor.execute("DELETE FROM dataset_files WHERE dataset_id = %s", (id,))
    cursor.execute("DELETE FROM dataset_manage WHERE id = %s", (id,))
    conn.commit()
    cursor.close()
    conn.close()
    return jsonify({'code': 0, 'message': '删除成功'})


# ============ 数据集文件上传接口 ============

def make_response_with_cors(data, status=200):
    """创建带CORS头的响应"""
    response = jsonify(data)
    response.headers['Access-Control-Allow-Origin'] = '*'
    response.headers['Access-Control-Allow-Methods'] = 'GET, POST, OPTIONS'
    response.headers['Access-Control-Allow-Headers'] = 'Content-Type'
    return response

@datasets_bp.route('/upload/<int:dataset_id>', methods=['POST', 'OPTIONS'])
def upload_dataset_file(dataset_id):
    """上传数据集文件"""
    # 处理 OPTIONS 预检请求
    if request.method == 'OPTIONS':
        response = make_response_with_cors({'status': 'ok'})
        return response

    # 检查数据集是否存在
    dataset = generic_get_by_id('dataset_manage', dataset_id)
    if not dataset:
        return make_response_with_cors({'code': 1, 'message': '数据集不存在'})

    # 确保上传目录存在（datasets根目录）
    os.makedirs(DATASET_FOLDER, exist_ok=True)

    uploaded_files = []
    errors = []

    if 'files' not in request.files:
        return make_response_with_cors({'code': 1, 'message': '没有文件被上传'})

    files = request.files.getlist('files')

    for file in files:
        if file.filename == '':
            continue

        if file and allowed_file(file.filename):
            filename = secure_filename(file.filename)
            # 添加时间戳和dataset_id防止文件名冲突，格式：timestamp_datasetId_filename
            timestamp = int(time.time() * 1000)
            new_filename = f"{timestamp}_{dataset_id}_{filename}"
            file_path = os.path.join(DATASET_FOLDER, new_filename)

            # 保存文件
            file.save(file_path)
            file_size = os.path.getsize(file_path)

            # 获取文件扩展名（安全处理无扩展名的情况）
            parts = filename.rsplit('.', 1)
            ext = parts[1].lower() if len(parts) > 1 else 'unknown'

            # 保存文件信息到数据库
            conn = get_db_connection()
            cursor = conn.cursor()
            cursor.execute(
                "INSERT INTO dataset_files (dataset_id, file_name, file_path, file_size, file_type) VALUES (%s, %s, %s, %s, %s)",
                (dataset_id, filename, file_path, file_size, ext)
            )
            conn.commit()
            cursor.close()
            conn.close()

            uploaded_files.append({
                'name': filename,
                'size': file_size,
                'size_formatted': format_file_size(file_size)
            })
        else:
            errors.append(f"{file.filename}: 文件类型不支持")

    # 如果有成功上传的文件，才更新数据集的文件数量和大小
    if uploaded_files:
        conn = get_db_connection()
        cursor = conn.cursor()
        cursor.execute("SELECT COUNT(*) as count, SUM(file_size) as total_size FROM dataset_files WHERE dataset_id = %s", (dataset_id,))
        result = cursor.fetchone()
        file_count = result['count'] or 0
        total_size = result['total_size'] or 0

        cursor.execute(
            "UPDATE dataset_manage SET file_count = %s, size = %s WHERE id = %s",
            (file_count, format_file_size(total_size), dataset_id)
        )
        conn.commit()
        cursor.close()
        conn.close()

    if errors:
        return make_response_with_cors({
            'code': 0,
            'message': f'部分文件上传成功，{len(errors)}个文件失败',
            'data': {
                'uploaded': uploaded_files,
                'errors': errors
            }
        })

    return make_response_with_cors({
        'code': 0,
        'message': f'成功上传 {len(uploaded_files)} 个文件',
        'data': {
            'uploaded': uploaded_files,
            'file_count': file_count
        }
    })


@datasets_bp.route('/<int:dataset_id>/files', methods=['GET'])
def get_dataset_files(dataset_id):
    """获取数据集文件列表"""
    conn = get_db_connection()
    cursor = conn.cursor()
    cursor.execute(
        "SELECT id, file_name, file_path, file_size, file_type, create_time FROM dataset_files WHERE dataset_id = %s ORDER BY create_time DESC",
        (dataset_id,)
    )
    files = cursor.fetchall()
    cursor.close()
    conn.close()

    # 格式化文件大小
    for f in files:
        f['file_size_formatted'] = format_file_size(f['file_size'])

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


@datasets_bp.route('/files/<int:file_id>', methods=['DELETE'])
def delete_dataset_file(file_id):
    """删除数据集文件"""
    conn = get_db_connection()
    cursor = conn.cursor()

    # 获取文件信息
    cursor.execute("SELECT dataset_id, file_path FROM dataset_files WHERE id = %s", (file_id,))
    file_info = cursor.fetchone()

    if not file_info:
        cursor.close()
        conn.close()
        return jsonify({'code': 1, 'message': '文件不存在'})

    # 删除物理文件
    file_path = file_info['file_path']
    if file_path and os.path.exists(file_path):
        try:
            os.remove(file_path)
        except Exception as e:
            print(f"删除文件失败: {file_path}, {e}")

    # 删除数据库记录
    cursor.execute("DELETE FROM dataset_files WHERE id = %s", (file_id,))

    # 更新数据集的文件数量和大小
    dataset_id = file_info['dataset_id']
    cursor.execute("SELECT COUNT(*) as count, SUM(file_size) as total_size FROM dataset_files WHERE dataset_id = %s", (dataset_id,))
    result = cursor.fetchone()
    file_count = result['count'] or 0
    total_size = result['total_size'] or 0

    cursor.execute(
        "UPDATE dataset_manage SET file_count = %s, size = %s WHERE id = %s",
        (file_count, format_file_size(total_size), dataset_id)
    )

    conn.commit()
    cursor.close()
    conn.close()

    return jsonify({'code': 0, 'message': '删除成功'})


# ============ 文件下载接口 ============

@datasets_bp.route('/download/<int:dataset_id>/<filename>', methods=['GET'])
def download_dataset_file(dataset_id, filename):
    """下载数据集文件"""
    # 文件直接存储在 DATASET_FOLDER 根目录下
    return send_from_directory(DATASET_FOLDER, filename, as_attachment=True)


@datasets_bp.route('/download/<int:dataset_id>', methods=['GET'])
def download_dataset_all(dataset_id):
    """下载数据集所有文件（ZIP打包）"""
    conn = get_db_connection()
    cursor = conn.cursor()

    # 获取数据集信息
    cursor.execute("SELECT name FROM dataset_manage WHERE id = %s", (dataset_id,))
    dataset = cursor.fetchone()

    if not dataset:
        cursor.close()
        conn.close()
        return jsonify({'code': 1, 'message': '数据集不存在'})

    # 获取所有文件
    cursor.execute(
        "SELECT id, file_name, file_path FROM dataset_files WHERE dataset_id = %s ORDER BY create_time DESC",
        (dataset_id,)
    )
    files = cursor.fetchall()
    cursor.close()
    conn.close()

    if not files:
        return jsonify({'code': 1, 'message': '数据集没有文件'})

    # 创建ZIP文件
    memory_file = io.BytesIO()
    with zipfile.ZipFile(memory_file, 'w', zipfile.ZIP_DEFLATED) as zf:
        for f in files:
            file_path = f.get('file_path')
            if file_path and os.path.exists(file_path):
                # 使用原始文件名
                zf.write(file_path, f['file_name'])

    memory_file.seek(0)

    # 发送ZIP文件
    zip_name = f"{dataset['name'] or 'dataset'}_{dataset_id}.zip"
    return Response(
        memory_file,
        mimetype='application/zip',
        headers={'Content-Disposition': f'attachment;filename={zip_name}'}
    )


# ============ 文件预览接口 ============

@datasets_bp.route('/preview/<int:file_id>', methods=['GET'])
def preview_dataset_file(file_id):
    """预览数据集文件内容（限100KB）"""
    conn = get_db_connection()
    cursor = conn.cursor()

    # 获取文件信息
    cursor.execute("SELECT id, file_name, file_path, file_type FROM dataset_files WHERE id = %s", (file_id,))
    file_info = cursor.fetchone()

    if not file_info:
        cursor.close()
        conn.close()
        return jsonify({'code': 1, 'message': '文件不存在'})

    file_path = file_info['file_path']

    if not file_path or not os.path.exists(file_path):
        cursor.close()
        conn.close()
        return jsonify({'code': 1, 'message': '文件不存在'})

    # 读取文件内容（限100KB）
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read(102400)  # 100KB
    except Exception as e:
        cursor.close()
        conn.close()
        return jsonify({'code': 1, 'message': f'读取文件失败: {str(e)}'})

    cursor.close()
    conn.close()

    return jsonify({
        'code': 0,
        'data': {
            'file_name': file_info['file_name'],
            'content': content
        }
    })
