"""
精调训练 API 路由
调用 llamafactory-cli 执行训练任务
"""
import os
import sys
import subprocess
import json
import threading
import time
import signal
import yaml
from flask import Blueprint, request, jsonify
import logging

# 添加项目根目录到路径
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.insert(0, PROJECT_ROOT)

# 加载配置
CONFIG_PATH = os.path.join(PROJECT_ROOT, 'config.yaml')
with open(CONFIG_PATH, 'r', encoding='utf-8') as f:
    CONFIG = yaml.safe_load(f)

logger = logging.getLogger(__name__)
train_logger = logging.getLogger('train')  # 专门的训练日志 logger，输出到 train.log

# 从配置获取训练日志路径
TRAINING_LOGS_DIR = CONFIG.get('training_logs_path', '/app/base/training_logs')

# 创建蓝图
fine_tune_bp = Blueprint('fine_tune', __name__, url_prefix='/api/fine-tune')


# 训练类型映射
TRAIN_TYPE_MAP = {
    'SFT': 'sft',
    'DPO': 'dpo',
    'CPT': 'cpt'
}

# 训练方法映射
FINETUNING_TYPE_MAP = {
    'lora': 'lora',
    'full': 'full'
}


@fine_tune_bp.route('/start', methods=['POST'])
def start_training():
    """启动 llamafactory 训练任务"""
    try:
        data = request.json
        train_logger.info(f"[TRAIN] ========== 开始训练任务 ==========")
        train_logger.info(f"[TRAIN] 收到启动训练请求: base_model={data.get('base_model')}, train_dataset_id={data.get('train_dataset_id')}")

        # 必填参数验证
        required_fields = ['base_model', 'template', 'train_dataset_id']
        for field in required_fields:
            if not data.get(field):
                return jsonify({'code': 1, 'message': f'缺少必要参数: {field}'})

        # 获取模型信息
        model_path = data.get('base_model')
        # 尝试转换为整数
        try:
            model_id = int(model_path) if str(model_path).isdigit() else None
        except (ValueError, TypeError):
            model_id = None

        if model_id:
            # 如果是 model_id，需要获取模型路径
            from .model_manage import get_db_connection
            conn = get_db_connection()
            cursor = conn.cursor()
            cursor.execute("SELECT id, name, path FROM model_manage WHERE id = %s", (model_id,))
            result = cursor.fetchone()
            conn.close()
            logger.info(f"模型查询结果: {result}")
            if result and result.get('path'):
                model_path = result['path']
                logger.info(f"从数据库获取的模型路径: {model_path}")
            else:
                return jsonify({'code': 1, 'message': '模型不存在或路径为空'})
        elif not model_path:
            return jsonify({'code': 1, 'message': f'模型路径为空'})

        train_logger.info(f"[TRAIN] 模型路径: {model_path}")

        # 设置工作目录和 llamafactory 目录
        work_dir = '/app/base'
        llamafactory_dir = '/app/base'

        # 数据集目录直接使用 /app/base/datasets（不再复制）
        datasets_dir = '/app/base/datasets'

        # 获取数据集名称（用于 --dataset 参数）
        dataset_key = None
        dataset_id = data.get('train_dataset_id')
        try:
            dataset_id_int = int(dataset_id) if str(dataset_id).isdigit() else None
        except (ValueError, TypeError):
            dataset_id_int = None

        if dataset_id_int:
            from .datasets import get_db_connection as get_dataset_conn
            conn = get_dataset_conn()
            cursor = conn.cursor()
            cursor.execute("SELECT dm.name FROM dataset_manage dm WHERE dm.id = %s", (dataset_id_int,))
            dataset_result = cursor.fetchone()
            conn.close()
            dataset_key = dataset_result['name'] if dataset_result else None
            train_logger.info(f"[TRAIN] 数据集名称: {dataset_key}")

        # 获取选中的 GPU 索引
        gpus = data.get('gpus', [])
        if gpus:
            gpu_ids = [gpu.get('id', '').replace('gpu', '') for gpu in gpus]
            gpu_ids = [g for g in gpu_ids if g.isdigit()]
            cuda_devices = ','.join(gpu_ids)
        else:
            cuda_devices = '0'

        # 设置环境变量
        env = os.environ.copy()
        env['CUDA_VISIBLE_DEVICES'] = cuda_devices
        env['TF_CPP_MIN_LOG_LEVEL'] = '2'  # 减少 TensorFlow 日志
        env['LLAMAFACTORY_DIR'] = '/app/base'  # 指定 llamafactory 根目录
        env['PYTHONUNBUFFERED'] = '1'  # 强制 Python 不缓冲输出，实时写入日志
        env['TRANSFORMERS_VERBOSITY'] = 'INFO'  # 设置 transformers 日志级别

        # 构建 llamafactory-cli 命令（传入数据集名称用于 --dataset 参数）
        cmd = build_train_command(data, model_path, dataset_key)
        cmd_str = ' '.join(cmd)
        train_logger.info(f"[TRAIN] 执行训练命令: {cmd_str}")

        # 在返回的命令中显示 GPU 配置
        cmd_str_with_gpu = f"CUDA_VISIBLE_DEVICES={cuda_devices} {cmd_str}"

        # 生成训练日志文件路径（存储在 logs 目录下的日期子目录中）
        from datetime import datetime
        today = datetime.now().strftime('%Y-%m-%d')
        now_str = datetime.now().strftime('%Y%m%d_%H%M%S')  # 时间戳用于排序
        task_id = data.get('task_id', 'unknown')
        task_name = data.get('name', 'unknown')
        # 工作目录设为 /app/base（而非 llamafactory 目录）
        work_dir = '/app/base'
        # 使用 logs 目录下的日期子目录
        training_logs_dir = os.path.join('/app/base/logs', today)
        os.makedirs(training_logs_dir, exist_ok=True)

        # 日志文件路径: logs/{日期}/{task_id}_{task_name}.log
        log_file = os.path.join(training_logs_dir, f'{task_id}_{task_name}.log')

        train_logger.info(f"[TRAIN] 启动训练进程...")

        # 用于存储实际进程 PID
        actual_pid = None
        final_log_path = log_file

        # 使用线程在后台运行训练进程
        def run_training():
            nonlocal actual_pid, final_log_path

            # 从 data 中获取 template 和 train_method（与 build_train_command 保持一致）
            template = data.get('template', 'default')
            train_method = data.get('train_method', 'lora')

            # 创建输出目录（如果不存在）
            # 路径格式: /app/base/saves/{train_method}/{output_model_name}
            output_model_name = data.get('output_model_name', template)
            output_model_name = f"{train_method}/{output_model_name}"
            if not output_model_name.startswith('/'):
                output_model_name = f"/app/base/saves/{output_model_name}"
            output_dir = output_model_name
            os.makedirs(output_dir, exist_ok=True)
            train_logger.info(f"[TRAIN] 输出目录: {output_dir}")
            train_logger.info(f"[TRAIN] 完整训练命令: {' '.join(cmd)}")

            with open(log_file, 'w', encoding='utf-8') as f:
                # 设置 cwd 为 /app，但通过 LLAMAFACTORY_DIR 环境变量指定 llamafactory 位置
                process = subprocess.Popen(
                    cmd,
                    cwd=work_dir,
                    stdout=f,
                    stderr=subprocess.STDOUT,
                    env=env
                )
                actual_pid = process.pid
                train_logger.info(f"[TRAIN] 训练进程 PID: {actual_pid}")
                train_logger.info(f"[TRAIN] 日志文件: {log_file}")

                # 更新数据库中的 PID（立即更新，方便停止任务）
                update_fine_tune_status(task_id, 'running', actual_pid)

                # 等待进程完成
                process.wait()
                train_logger.info(f"[TRAIN] 训练进程已结束，退出码: {process.returncode}")

                # 更新任务状态
                final_status = 'completed' if process.returncode == 0 else 'failed'
                update_fine_tune_status(task_id, final_status, actual_pid)

        # 启动后台线程
        training_thread = threading.Thread(target=run_training, daemon=True)
        training_thread.start()

        # 等待 PID 并更新到数据库
        for i in range(10):  # 最多等待1秒
            time.sleep(0.1)
            if actual_pid:
                break

        # 立即返回，不等待进程完成
        train_logger.info(f"[TRAIN] 训练任务已在后台启动，PID: {actual_pid}")

        train_logger.info(f"[TRAIN] 训练日志输出到: {log_file}")

        return jsonify({
            'code': 0,
            'message': f'训练任务已启动 (GPU: {cuda_devices})',
            'data': {
                'task_id': task_id,
                'pid': actual_pid,
                'gpu_ids': cuda_devices,
                'command': cmd_str_with_gpu,
                'log_file': log_file,
                'training_logs_dir': training_logs_dir
            }
        })

    except Exception as e:
        import traceback
        train_logger.error(f"[TRAIN] 启动训练任务失败: {e}")
        train_logger.error(f"[TRAIN] 详细错误: {traceback.format_exc()}")
        return jsonify({'code': 1, 'message': str(e)})


def build_train_command(data, model_path, dataset_name=None):
    """构建 llamafactory-cli train 命令"""
    # llamafactory-cli 路径（已在系统 PATH 中）
    cmd = ['llamafactory-cli', 'train']

    # 训练阶段
    train_type = data.get('train_type', 'SFT')
    cmd.extend(['--stage', TRAIN_TYPE_MAP.get(train_type, 'sft')])
    cmd.append('--do_train')

    # 模型路径
    cmd.extend(['--model_name_or_path', model_path])

    # 数据集 - 使用数据集名称（dataset_manage.name），不是实际文件名
    if dataset_name:
        cmd.extend(['--dataset', dataset_name])
        train_logger.info(f"[TRAIN] 使用数据集名称: {dataset_name}")
    else:
        # 回退到原有逻辑
        dataset_id = data.get('train_dataset_id')
        try:
            dataset_id_int = int(dataset_id) if str(dataset_id).isdigit() else None
        except (ValueError, TypeError):
            dataset_id_int = None

        if dataset_id_int:
            dataset_name = get_dataset_name(dataset_id_int)
            train_logger.info(f"[TRAIN] 从数据库获取的数据集名称: {dataset_name}")
        else:
            dataset_name = dataset_id
        cmd.extend(['--dataset', dataset_name])

    # 数据集目录
    cmd.extend(['--dataset_dir', './datasets'])  # llamafactory 工作目录下的 datasets 目录

    # 模板
    template = data.get('template')
    cmd.extend(['--template', template])

    # 训练方法
    train_method = data.get('train_method', 'lora')
    cmd.extend(['--finetuning_type', FINETUNING_TYPE_MAP.get(train_method, 'lora')])

    # 输出目录（确保是绝对路径）
    # 路径格式: /app/base/saves/{train_method}/{output_model_name}
    output_model_name = data.get('output_model_name', template)
    output_model_name = f"{train_method}/{output_model_name}"
    if not output_model_name.startswith('/'):
        output_model_name = f"/app/base/saves/{output_model_name}"
    output_dir = output_model_name
    cmd.extend(['--output_dir', output_dir])

    # 常用参数
    cmd.extend([
        '--overwrite_cache',
        '--overwrite_output_dir',
        '--cutoff_len', str(data.get('max_length', 512)),
        '--preprocessing_num_workers', '16',
        '--per_device_train_batch_size', str(data.get('batch_size', 1)),
        '--per_device_eval_batch_size', '1',
        '--gradient_accumulation_steps', str(data.get('gradient_accumulation_steps', 8)),
        '--lr_scheduler_type', data.get('lr_scheduler_type', 'cosine'),
        '--logging_steps', '5',
        '--warmup_steps', str(data.get('warmup_steps', 20)),
        '--save_steps', str(data.get('save_steps', 100)),
        '--log_level', 'info',  # 设置日志级别为 info
        '--log_level_replica', 'info',  # 设置副本日志级别
    ])

    # 学习率
    cmd.extend(['--learning_rate', str(data.get('learning_rate', 0.0001))])

    # 训练轮数
    cmd.extend(['--num_train_epochs', str(data.get('n_epochs', 1.0))])

    # 验证集比例
    val_ratio = data.get('valid_ratio', 0)
    if val_ratio > 0:
        cmd.extend(['--val_size', str(val_ratio / 100)])

    # 最大样本数
    if data.get('max_samples'):
        cmd.extend(['--max_samples', str(data.get('max_samples'))])

    # 启用 TensorBoard 日志（用于可视化训练曲线）
    cmd.append('--plot_loss')

    # 其他选项

    if data.get('fp16'):
        cmd.append('--fp16')

    if data.get('load_best_model_at_end'):
        cmd.append('--load_best_model_at_end')

    return cmd


def get_dataset_name(dataset_id):
    """根据数据集 ID 获取数据集名称"""
    try:
        from .datasets import get_db_connection
        conn = get_db_connection()
        cursor = conn.cursor()
        cursor.execute("SELECT id, name FROM dataset_manage WHERE id = %s", (dataset_id,))
        result = cursor.fetchone()
        conn.close()
        logger.info(f"数据集查询结果: {result}")
        if result and result.get('name'):
            return result['name']
        logger.warning(f"未找到数据集 ID={dataset_id}，使用默认值")
        return 'default'
    except Exception as e:
        logger.error(f"查询数据集失败: {e}")
        return 'default'


def update_fine_tune_status(task_id, status, pid=None):
    """更新训练任务状态"""
    try:
        from .model_manage import get_db_connection
        conn = get_db_connection()
        cursor = conn.cursor()

        if status == 'running' and pid:
            cursor.execute(
                "UPDATE fine_tune SET status = %s, process_id = %s WHERE id = %s",
                (status, pid, task_id)
            )
        else:
            cursor.execute(
                "UPDATE fine_tune SET status = %s WHERE id = %s",
                (status, task_id)
            )

        conn.commit()
        conn.close()
    except Exception as e:
        logger.error(f"更新任务状态失败: {e}")


@fine_tune_bp.route('/stop/<int:task_id>', methods=['POST'])
def stop_training(task_id):
    """停止训练任务"""
    try:
        from .model_manage import get_db_connection

        # 获取进程 ID
        conn = get_db_connection()
        cursor = conn.cursor()
        cursor.execute("SELECT process_id FROM fine_tune WHERE id = %s", (task_id,))
        result = cursor.fetchone()
        conn.close()

        if result and result.get('process_id'):
            pid = result['process_id']
            try:
                # 尝试终止进程
                import signal
                os.kill(pid, signal.SIGTERM)
                logger.info(f"已终止训练进程 PID: {pid}")
            except ProcessLookupError:
                logger.warning(f"进程 {pid} 不存在")
            except PermissionError:
                logger.error(f"没有权限终止进程 {pid}")

        # 更新状态
        update_fine_tune_status(task_id, 'stopped')

        return jsonify({'code': 0, 'message': '训练任务已停止'})

    except Exception as e:
        logger.error(f"停止训练任务失败: {e}")
        return jsonify({'code': 1, 'message': str(e)})


@fine_tune_bp.route('/<int:task_id>', methods=['DELETE'])
def delete_training_task(task_id):
    """删除训练任务及对应的日志文件"""
    try:
        from .model_manage import get_db_connection

        # 获取任务信息（用于删除日志文件）
        conn = get_db_connection()
        cursor = conn.cursor()
        # 尝试获取所有字段，如果tensorboard_log_dir不存在会报错
        try:
            cursor.execute("SELECT name, process_id, tensorboard_log_dir FROM fine_tune WHERE id = %s", (task_id,))
        except:
            # 如果列不存在，只获取基本字段
            cursor.execute("SELECT name, process_id FROM fine_tune WHERE id = %s", (task_id,))
        task_result = cursor.fetchone()
        conn.close()

        if not task_result:
            return jsonify({'code': 1, 'message': '任务不存在'})

        task_name = task_result.get('name', 'unknown')
        tensorboard_log_dir = task_result.get('tensorboard_log_dir', '') if 'tensorboard_log_dir' in task_result else ''

        # 删除日志文件 (logs/{日期}/{task_id}_{task_name}.log)
        try:
            from datetime import datetime
            today = datetime.now().strftime('%Y-%m-%d')

            # 可能的日志文件路径
            log_paths = [
                f'/app/base/logs/{today}/{task_id}_{task_name}.log',
                f'/app/base/logs/{task_id}_{task_name}.log',
            ]

            for log_path in log_paths:
                if os.path.exists(log_path):
                    os.remove(log_path)
                    logger.info(f"已删除日志文件: {log_path}")
        except Exception as log_err:
            logger.warning(f"删除日志文件失败: {log_err}")

        # 删除 TensorBoard 进程（如果存在）
        global tensorboard_process
        if tensorboard_process and tensorboard_process.poll() is None:
            try:
                import signal
                os.killpg(os.getpgid(tensorboard_process.pid), signal.SIGTERM)
                tensorboard_process = None
                logger.info(f"已停止 TensorBoard 进程")
            except Exception as tb_err:
                logger.warning(f"停止 TensorBoard 失败: {tb_err}")

        # 删除数据库中的任务记录
        conn = get_db_connection()
        cursor = conn.cursor()
        cursor.execute("DELETE FROM fine_tune WHERE id = %s", (task_id,))
        conn.commit()
        conn.close()

        logger.info(f"已删除训练任务 {task_id}: {task_name}")
        return jsonify({'code': 0, 'message': '删除成功'})

    except Exception as e:
        logger.error(f"删除训练任务失败: {e}")
        return jsonify({'code': 1, 'message': str(e)})


@fine_tune_bp.route('/status/<int:task_id>', methods=['GET'])
def get_training_status(task_id):
    """获取训练任务状态"""
    try:
        from .model_manage import get_db_connection

        conn = get_db_connection()
        cursor = conn.cursor()
        cursor.execute(
            "SELECT id, name, status, progress, process_id FROM fine_tune WHERE id = %s",
            (task_id,)
        )
        result = cursor.fetchone()
        conn.close()

        if result:
            # 检查 PID 是否仍在运行
            actual_status = result['status']
            pid = result.get('process_id')
            if pid and actual_status == 'running':
                try:
                    # 检查进程是否存在
                    os.kill(pid, 0)
                    # 进程仍在运行
                    actual_status = 'running'
                except (OSError, ProcessLookupError):
                    # 进程已结束，尝试更新状态
                    actual_status = 'completed'  # 假设完成（实际可能失败）

            return jsonify({
                'code': 0,
                'data': {
                    'task_id': result['id'],
                    'name': result['name'],
                    'status': actual_status,
                    'progress': result['progress'],
                    'pid': result.get('process_id')
                }
            })
        else:
            return jsonify({'code': 1, 'message': '任务不存在'})

    except Exception as e:
        logger.error(f"获取任务状态失败: {e}")
        return jsonify({'code': 1, 'message': str(e)})


@fine_tune_bp.route('/check-pid/<int:pid>', methods=['GET'])
def check_pid_status(pid):
    """检查 PID 是否仍在运行"""
    try:
        if pid <= 0:
            return jsonify({
                'code': 0,
                'data': {
                    'exists': False,
                    'message': '无效的 PID'
                }
            })

        try:
            # 发送信号 0 来检查进程是否存在（不会实际终止进程）
            os.kill(pid, 0)
            return jsonify({
                'code': 0,
                'data': {
                    'exists': True,
                    'message': '进程仍在运行'
                }
            })
        except (OSError, ProcessLookupError):
            # 进程不存在
            return jsonify({
                'code': 0,
                'data': {
                    'exists': False,
                    'message': '进程已结束'
                }
            })
    except Exception as e:
        logger.error(f"检查 PID 状态失败: {e}")
        return jsonify({
            'code': 0,
            'data': {
                'exists': False,
                'message': f'检查失败: {str(e)}'
            }
        })


@fine_tune_bp.route('/log/<int:task_id>', methods=['GET'])
def get_training_log(task_id):
    """获取训练任务日志内容（支持实时读取）"""
    try:
        from .model_manage import get_db_connection

        # 获取任务信息和进程ID
        conn = get_db_connection()
        cursor = conn.cursor()
        cursor.execute(
            "SELECT name, process_id, status FROM fine_tune WHERE id = %s",
            (task_id,)
        )
        result = cursor.fetchone()
        conn.close()

        if not result:
            return jsonify({'code': 1, 'message': '任务不存在'})

        process_id = result.get('process_id')
        task_name = result['name']
        status = result['status']

        if not process_id:
            return jsonify({'code': 1, 'message': '任务尚未启动'})

        # 构建日志文件路径 - 新格式: logs/{日期}/{task_id}_{task_name}.log
        from datetime import datetime
        today = datetime.now().strftime('%Y-%m-%d')
        training_logs_dir = os.path.join('/app/base/logs', today)

        # 查找日志文件 (新格式: {task_id}_{task_name}.log)
        log_file = os.path.join(training_logs_dir, f'{task_id}_{task_name}.log')

        if not os.path.exists(log_file):
            # 如果没找到，返回空日志
            return jsonify({
                'code': 0,
                'data': {
                    'content': '',
                    'status': status,
                    'message': '日志文件尚未创建'
                }
            })

        # 读取日志文件内容
        try:
            with open(log_file, 'r', encoding='utf-8', errors='ignore') as f:
                content = f.read()
            return jsonify({
                'code': 0,
                'data': {
                    'content': content,
                    'status': status,
                    'log_file': log_file
                }
            })
        except Exception as e:
            return jsonify({
                'code': 0,
                'data': {
                    'content': '',
                    'status': status,
                    'message': f'读取日志失败: {str(e)}'
                }
            })

    except Exception as e:
        logger.error(f"获取训练日志失败: {e}")
        return jsonify({'code': 1, 'message': str(e)})


import re


@fine_tune_bp.route('/progress/<int:task_id>', methods=['GET'])
def get_training_progress(task_id):
    """获取训练任务进度（从日志中解析 llamafactory 的进度信息）"""
    try:
        from .model_manage import get_db_connection

        # 获取任务信息和进程ID
        conn = get_db_connection()
        cursor = conn.cursor()
        cursor.execute(
            "SELECT name, process_id, status FROM fine_tune WHERE id = %s",
            (task_id,)
        )
        result = cursor.fetchone()
        conn.close()

        if not result:
            return jsonify({'code': 1, 'message': '任务不存在'})

        process_id = result.get('process_id')
        task_name = result['name']
        status = result['status']

        if not process_id:
            return jsonify({
                'code': 0,
                'data': {
                    'progress': 0,
                    'step': '',
                    'eta': '',
                    'speed': '',
                    'status': status,
                    'message': '任务尚未启动'
                }
            })

        # 构建日志文件路径 - 新格式: logs/{日期}/{task_id}_{task_name}.log
        from datetime import datetime
        today = datetime.now().strftime('%Y-%m-%d')
        training_logs_dir = os.path.join('/app/base/logs', today)

        # 查找日志文件 (新格式: {task_id}_{task_name}.log)
        log_file = os.path.join(training_logs_dir, f'{task_id}_{task_name}.log')

        # TensorBoard 日志目录（使用默认值）
        tensorboard_log_dir = '/app/base/saves'

        if not os.path.exists(log_file):
            return jsonify({
                'code': 0,
                'data': {
                    'step': '',
                    'elapsed': '',
                    'eta': '',
                    'speed': '',
                    'status': status,
                    'message': '日志文件尚未创建',
                    'tensorboard_url': ''
                }
            })

        # 读取日志文件最后部分，解析进度信息
        try:
            with open(log_file, 'r', encoding='utf-8', errors='ignore') as f:
                # 读取最后 10KB 内容
                f.seek(0, 2)  # 跳到文件末尾
                file_size = f.tell()
                read_size = min(10240, file_size)
                f.seek(max(0, file_size - read_size))
                content = f.read()

            # 匹配 llamafactory 进度格式:  52%|█████▏    | 17/33 [02:16<02:08,  8.04s/it]
            progress_pattern = r'\s*(\d+)%\|[█░▌▋▒█▏▎▏▐▀■□▪▫‣▶➜➡→]+\s*\|\s*(\d+)/(\d+)\s+\[(\d+):(\d+)<(\d+):(\d+),\s*([\d.]+)s/it\]'
            match = re.search(progress_pattern, content)

            step_info = ''
            elapsed = ''
            eta = ''
            speed = ''
            message = '等待训练开始'

            if match:
                current_step = int(match.group(2))
                total_steps = int(match.group(3))
                elapsed_min = int(match.group(4))
                elapsed_sec = int(match.group(5))
                eta_min = int(match.group(6))
                eta_sec = int(match.group(7))
                speed_val = float(match.group(8))

                step_info = f'{current_step}/{total_steps}'
                elapsed = f'{elapsed_min:02d}:{elapsed_sec:02d}'
                eta = f'{eta_min:02d}:{eta_sec:02d}'
                speed = f'{speed_val}s/it'
                message = '训练进行中'

            return jsonify({
                'code': 0,
                'data': {
                    'step': step_info,
                    'elapsed': elapsed,
                    'eta': eta,
                    'speed': speed,
                    'status': status,
                    'message': message,
                    'tensorboard_log_dir': tensorboard_log_dir,
                    'tensorboard_url': ''
                }
            })

        except Exception as e:
            return jsonify({
                'code': 0,
                'data': {
                    'step': '',
                    'elapsed': '',
                    'eta': '',
                    'speed': '',
                    'status': status,
                    'message': f'读取进度失败: {str(e)}',
                    'tensorboard_log_dir': tensorboard_log_dir,
                    'tensorboard_url': ''
                }
            })

    except Exception as e:
        logger.error(f"获取训练进度失败: {e}")
        return jsonify({'code': 1, 'message': str(e)})


def get_db_connection():
    """获取数据库连接"""
    import pymysql
    import yaml

    PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    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
    )


@fine_tune_bp.route('/check-name', methods=['GET'])
def check_task_name():
    """检查任务名称是否重复"""
    try:
        name = request.args.get('name', '').strip()
        if not name:
            return jsonify({'code': 1, 'message': '任务名称不能为空'})

        # 验证任务名称格式：只能包含英文、数字、下划线
        import re
        if not re.match(r'^[a-zA-Z0-9_]+$', name):
            return jsonify({'code': 1, 'message': '任务名称只能包含英文、数字和下划线'})

        conn = get_db_connection()
        cursor = conn.cursor()
        cursor.execute("SELECT id FROM fine_tune WHERE name = %s", (name,))
        result = cursor.fetchone()
        conn.close()

        if result:
            return jsonify({
                'code': 0,
                'data': {
                    'exists': True,
                    'message': '任务名称已存在'
                }
            })
        else:
            return jsonify({
                'code': 0,
                'data': {
                    'exists': False,
                    'message': '任务名称可用'
                }
            })

    except Exception as e:
        logger.error(f"检查任务名称失败: {e}")
        return jsonify({'code': 1, 'message': str(e)})


# TensorBoard 服务进程
tensorboard_process = None


@fine_tune_bp.route('/tensorboard/start', methods=['POST'])
def start_tensorboard():
    """启动 TensorBoard 服务"""
    global tensorboard_process
    try:
        import subprocess
        import os

        # 检查是否已有进程在运行
        if tensorboard_process and tensorboard_process.poll() is None:
            return jsonify({
                'code': 0,
                'data': {
                    'url': 'http://10.10.10.177:6006',
                    'status': 'already_running',
                    'message': 'TensorBoard 服务已运行'
                }
            })

        # 获取日志目录
        log_dir = '/app/base/saves'

        # 检查目录是否存在
        if not os.path.exists(log_dir):
            return jsonify({'code': 1, 'message': f'日志目录不存在: {log_dir}'})

        # 启动 TensorBoard（后台运行）
        cmd = ['tensorboard', '--logdir', log_dir, '--port', '6006', '--bind_all']
        tensorboard_process = subprocess.Popen(
            cmd,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            preexec_fn=os.setsid
        )

        logger.info(f"TensorBoard 服务已启动: {cmd}")

        return jsonify({
            'code': 0,
            'data': {
                'url': 'http://10.10.10.177:6006',
                'status': 'started',
                'message': 'TensorBoard 服务已启动'
            }
        })

    except Exception as e:
        logger.error(f"启动 TensorBoard 失败: {e}")
        return jsonify({'code': 1, 'message': str(e)})


@fine_tune_bp.route('/tensorboard/stop', methods=['POST'])
def stop_tensorboard():
    """停止 TensorBoard 服务"""
    global tensorboard_process
    try:
        import subprocess
        import signal

        if tensorboard_process and tensorboard_process.poll() is None:
            # 使用 os.killpg 终止进程组
            try:
                os.killpg(os.getpgid(tensorboard_process.pid), signal.SIGTERM)
            except Exception:
                pass
            tensorboard_process = None
            logger.info("TensorBoard 服务已停止")

        return jsonify({
            'code': 0,
            'data': {
                'status': 'stopped',
                'message': 'TensorBoard 服务已停止'
            }
        })

    except Exception as e:
        logger.error(f"停止 TensorBoard 失败: {e}")
        return jsonify({'code': 1, 'message': str(e)})
