"""
AI 补全训练工具

含 generate_table_by_description 的轮询配套 get_ai_training_detail。
训练状态 trainingStatus: 0 待训练 / 1 训练中 / 2 已完成 / 3 失败
"""

from ._base import register_tool, ToolDef


@register_tool("list_ai_trainings")
class ListAiTrainingsTool(ToolDef):
    name = "list_ai_trainings"
    description = "获取 AI 补全训练任务列表"
    input_schema = {
        "type": "object",
        "properties": {
            "datasourceId": {"type": "string", "description": "数据源 ID（可选过滤）"},
            "pageNum": {"type": "integer", "default": 1, "description": "页码"},
            "pageSize": {"type": "integer", "default": 20, "description": "每页数量"},
        },
        "required": [],
    }

    async def execute(self, args: dict) -> dict:
        params = {k: v for k, v in args.items() if v is not None}
        return await self.client.get("/ai/training/list", params=params)


@register_tool("create_ai_training_by_selected")
class CreateAiTrainingBySelectedTool(ToolDef):
    name = "create_ai_training_by_selected"
    description = "按选中的表创建 AI 补全训练任务"
    input_schema = {
        "type": "object",
        "properties": {
            "datasourceId": {"type": "string", "description": "数据源 ID"},
            "tableIds": {
                "type": "array",
                "description": "要训练的表 ID 数组",
                "items": {"type": "string"},
            },
        },
        "required": ["tableIds"],
    }

    async def execute(self, args: dict) -> dict:
        return await self.client.post("/ai/training/createBySelected", json_data=args)


@register_tool("get_ai_training_detail")
class GetAiTrainingDetailTool(ToolDef):
    name = "get_ai_training_detail"
    description = (
        "查询 AI 训练任务详情（轮询用）。"
        "trainingStatus: 0 待训练 / 1 训练中 / 2 已完成 / 3 失败。"
        "generate_table_by_description 返回 taskId 后，用本工具轮询直到 status=2 取结果"
    )
    input_schema = {
        "type": "object",
        "properties": {
            "taskId": {"type": "string", "description": "训练任务 ID（generate_table 返回的 taskId）"},
        },
        "required": ["taskId"],
    }

    async def execute(self, args: dict) -> dict:
        return await self.client.get(f"/ai/training/{args['taskId']}")
