from fastapi import APIRouter, HTTPException
from pydantic import BaseModel

from src.database.db_instance import db_manager
from src.core.task_manager import task_manager
import logging

router = APIRouter(prefix="/models", tags=["models"])


class CreateModelRequest(BaseModel):
    """
    拆分id
    """
    split_id: str
    """
    拆分名称
    """
    split_name: str
    """
    模型id
    """
    model_id: str
    """
    模型配置路径
    """
    model_config_path: str


@router.post("/create")
async def create_model(request: CreateModelRequest):
    """
    创建模型信息
    """
    try:
        success = db_manager.insert_model_info(
            model_id=request.model_id,
            model_name=request.model_id,
            split_id=request.split_id,
            split_name=request.split_name,
            model_info=request.model_config_path,
            state=0,
            assess_state=0,
            export_state=0,
        )

        if success:
            return {
                "status": "success",
                "message": "模型信息创建成功",
                "split_id": request.split_id,
            }
        else:
            raise HTTPException(status_code=500, detail="模型信息创建失败")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/split/{model_id}")
async def get_model_info_by_model_id(model_id: str):
    """
    获取模型信息
    """
    try:
        model_info_list = db_manager.get_model_info_by_model_id(model_id)
        if not model_info_list or len(model_info_list) == 0:
            raise HTTPException(status_code=404, detail="未找到对应的模型信息")
        return {"model_info": model_info_list}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/records/{model_id}")
async def get_model_records_by_model_id(model_id: str):
    """
    获取模型训练记录
    """
    try:
        model_records = db_manager.get_model_records_by_model_id(model_id)
        return {"model_records": model_records, "total": len(model_records)}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/latest/{model_id}")
async def get_latest_model_data(model_id: str):
    """
    获取模型需要要的最新数据，包括损失模型和评估指标
    """
    try:
        model_info_list = db_manager.get_model_info_by_model_id(model_id)

        if not model_info_list or len(model_info_list) == 0:
            raise HTTPException(status_code=404, detail="未找到对应的模型信息")

        # 获取最新的模型信息记录
        model_info = model_info_list[0]

        lossModel = db_manager.get_latest_model_record_by_type_and_model_id(
            1, model_id=model_id
        )
        mertis_record = db_manager.get_latest_model_record_by_type_and_model_id(
            2, model_id=model_id
        )
        
        # 如果有评估指标记录，解析metrics字段
        mertis = None
        if mertis_record and mertis_record.get("metrics"):
            try:
                import json
                # 将单引号替换为双引号以便JSON解析
                metrics_str = mertis_record["metrics"].replace("'", '"')
                mertis = json.loads(metrics_str)
                logging.info(f"解析后的评估指标: {mertis}")
            except Exception as e:
                logging.error(f"解析评估指标失败: {e}")
                # 如果解析失败，返回原始记录
                mertis = mertis_record
        
        logging.info(f"获取到的损失模型: {lossModel}, 评估指标: {mertis}")
        
        return {"lossModel": lossModel, "mertis": mertis, "model_info": model_info}
    except HTTPException as he:
        raise he
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


class UpdateRemarkRequest(BaseModel):
    """
    更新备注请求
    """
    model_id: str
    remark: str


@router.post("/remark")
async def update_model_remark(request: UpdateRemarkRequest):
    """ 
    更新模型备注信息
    """
    try:
        success = db_manager.update_model_info(request.model_id, remark=request.remark)

        if success:
            return {
                "status": "success",
                "message": "模型备注更新成功",
                "model_id": request.model_id,
            }
        else:
            raise HTTPException(status_code=500, detail="模型备注更新失败")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/stop/{model_id}")
async def stop_training(model_id: str):
    """
    停止模型训练
    """
    try:
        # 先获取模型信息中的task_id
        model_info_list = db_manager.get_model_info_by_model_id(model_id)
        if not model_info_list or len(model_info_list) == 0:
            raise HTTPException(status_code=404, detail="未找到对应的模型信息")

        # 获取最新的模型信息记录
        model_info = model_info_list[0]
        task_id = model_info.get('task_id')
        if task_id:
            # 停止任务管理器中的任务
            task_manager.stop_task(task_id)
            logging.info(f"Task {task_id} stopped via task_manager")

        # 更新数据库状态
        success = db_manager.update_model_info(model_id, state=5)
        if success:
            return {
                "status": "success",
                "message": "模型训练已停止",
                "model_id": model_id,
            }
        else:
            raise HTTPException(status_code=500, detail="停止训练失败")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/stop/task/{task_id}")
async def stop_training_by_task(task_id: str):
    """
    根据任务ID停止模型训练
    """
    try:
        # 停止任务管理器中的任务
        task_manager.stop_task(task_id)
        logging.info(f"Task {task_id} stopped via task_manager")

        # 更新数据库状态
        success = db_manager.stop_training_by_task_id(task_id)
        if success:
            return {
                "status": "success",
                "message": "模型训练已停止",
                "task_id": task_id,
            }
        else:
            raise HTTPException(status_code=500, detail="停止训练失败")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/delete/{model_id}")
async def delete_model(model_id: str):
    """"
    删除模型
    """
    try:
        success = db_manager.delete_model_info(model_id)
        if success:
            return {
                "status": "success",
                "message": "模型删除成功",
                "model_id": model_id,
            }
        else:
            raise HTTPException(status_code=500, detail="模型删除失败")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))