"""
数据导入工具 (工具 30-31)
"""

import base64
import io

from ._base import register_tool, ToolDef


@register_tool("preview_import_data")
class PreviewImportDataTool(ToolDef):
    name = "preview_import_data"
    description = "上传 Excel 文件，AI 智能识别并预览表结构/数据"
    input_schema = {
        "type": "object",
        "properties": {
            "connectionId": {"type": "string", "description": "数据源 ID"},
            "target": {"type": "string", "enum": ["prod", "test"], "default": "test", "description": "环境，默认 test"},
            "file_base64": {"type": "string", "description": "Excel 文件 base64 编码（.xlsx/.xls, <500KB）"},
            "file_name": {"type": "string", "description": "文件名（如 data.xlsx）"},
        },
        "required": ["connectionId", "file_base64"],
    }

    async def execute(self, args: dict) -> dict:
        args = dict(args)
        connection_id = args.pop("connectionId")
        target = args.pop("target", "test")
        file_base64 = args.pop("file_base64")
        file_name = args.pop("file_name", "import.xlsx")
        
        # 解码 base64 文件
        file_content = base64.b64decode(file_base64)
        
        # 构建文件上传
        files = {
            "file": (file_name, io.BytesIO(file_content), "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"),
        }
        
        return self.client.upload(
            f"/api/datasource/connection/{connection_id}/import_document/preview",
            files=files,
            params={"target": target},
        )


@register_tool("confirm_import_data")
class ConfirmImportDataTool(ToolDef):
    name = "confirm_import_data"
    description = "确认导入 AI 识别后的数据"
    input_schema = {
        "type": "object",
        "properties": {
            "connectionId": {"type": "string", "description": "数据源 ID"},
            "target": {"type": "string", "enum": ["prod", "test"], "description": "环境"},
            "data": {"type": "object", "description": "导入数据（含 tableStructure + allData）"},
        },
        "required": ["connectionId", "data"],
    }

    async def execute(self, args: dict) -> dict:
        args = dict(args)
        connection_id = args.pop("connectionId")
        target = args.pop("target", "test")
        data = args.pop("data")
        
        return self.client.post(
            f"/api/datasource/connection/{connection_id}/import_document/confirm",
            json_data=data,
            params={"target": target},
        )
