from fastapi import APIRouter, HTTPException, Query

from app.schemas.train_management import CreateTaskRequest, SaveModelRequest, UpdateModelRequest
from app.services.train_service import TrainService
from app.utils.response import success_response

router = APIRouter()
service = TrainService()


# ── packages ──────────────────────────────────────────────────────────────────

@router.get('/packages')
def list_packages(
    category_id: str = Query(default=''),
    name: str = Query(default=''),
):
    return success_response(data=service.list_packages(category_id=category_id, name=name.strip()))


# ── tasks ─────────────────────────────────────────────────────────────────────

@router.get('/tasks')
def list_tasks():
    return success_response(data=service.list_tasks())


@router.get('/tasks/{task_id}')
def get_task(task_id: int):
    try:
        return success_response(data=service.get_task(task_id))
    except ValueError as error:
        raise HTTPException(status_code=404, detail=str(error)) from error


@router.post('/tasks')
def create_task(request: CreateTaskRequest):
    try:
        task = service.create_task(
            model_name=request.model_name.strip(),
            train_package_id=request.train_package_id,
            val_package_id=request.val_package_id,
            params=request.params.model_dump(),
        )
        return success_response(data=task, message='训练任务已启动')
    except ValueError as error:
        raise HTTPException(status_code=400, detail=str(error)) from error


@router.post('/tasks/{task_id}/cancel')
def cancel_task(task_id: int):
    try:
        service.cancel_task(task_id)
        return success_response(data=True, message='取消请求已发送')
    except ValueError as error:
        raise HTTPException(status_code=400, detail=str(error)) from error


@router.post('/tasks/{task_id}/restart')
def restart_task(task_id: int):
    try:
        task = service.restart_task(task_id)
        return success_response(data=task, message='重新训练已启动')
    except ValueError as error:
        raise HTTPException(status_code=400, detail=str(error)) from error


@router.delete('/tasks/{task_id}')
def delete_task(task_id: int):
    try:
        service.delete_task(task_id)
        return success_response(data=True, message='任务已删除')
    except ValueError as error:
        raise HTTPException(status_code=400, detail=str(error)) from error


@router.post('/tasks/{task_id}/save')
def save_model(task_id: int, request: SaveModelRequest):
    try:
        model = service.save_model(task_id, model_name=request.model_name, description=request.description)
        return success_response(data=model, message='模型已保存')
    except ValueError as error:
        raise HTTPException(status_code=400, detail=str(error)) from error


# ── saved models ──────────────────────────────────────────────────────────────

@router.get('/models')
def list_models():
    return success_response(data=service.list_saved_models())


@router.delete('/models/{model_id}')
def delete_model(model_id: int):
    try:
        service.delete_saved_model(model_id)
        return success_response(data=True, message='模型已删除')
    except ValueError as error:
        raise HTTPException(status_code=400, detail=str(error)) from error

@router.patch('/models/{model_id}')
def update_model(model_id: int, request: UpdateModelRequest):
    try:
        model = service.update_saved_model(
            model_id,
            model_name=request.model_name,
            description=request.description,
        )
        return success_response(data=model, message='已更新')
    except ValueError as error:
        raise HTTPException(status_code=404, detail=str(error)) from error