from fastapi import APIRouter, HTTPException, Query
from pydantic import BaseModel, Field

from app.services.train_service import TrainService
from app.utils.response import success_response

router = APIRouter()
service = TrainService()


# ── request schemas ───────────────────────────────────────────────────────────

class LSTMParams(BaseModel):
    seq_len: int = Field(default=20, ge=5, le=500)
    hidden_size: int = Field(default=64, ge=8, le=1024)
    num_layers: int = Field(default=2, ge=1, le=8)
    epochs: int = Field(default=50, ge=1, le=2000)
    batch_size: int = Field(default=32, ge=1, le=512)
    learning_rate: float = Field(default=0.001, gt=0, le=1)


class CreateTaskRequest(BaseModel):
    model_name: str = Field(min_length=1, max_length=255)
    train_package_id: int
    val_package_id: int
    params: LSTMParams = Field(default_factory=LSTMParams)


class SaveModelRequest(BaseModel):
    model_name: str | None = Field(default=None, max_length=255)
    description: str | None = None


class UpdateModelRequest(BaseModel):
    model_name: str | None = Field(default=None, max_length=255)
    description: str | None = None


# ── 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