from fastapi import APIRouter, HTTPException
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)
    train_ratio: float = Field(default=0.8, ge=0.5, le=0.99)


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


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

@router.get('/packages')
def list_packages():
    return success_response(data=service.list_packages())


# ── 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(),
            package_id=request.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):
    try:
        model = service.save_model(task_id)
        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
