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

import io
import mimetypes
import os
from urllib.parse import urlparse, unquote

from ._base import register_tool, ToolDef


@register_tool("preview_import_data")
class PreviewImportDataTool(ToolDef):
    name = "preview_import_data"
    description = "从 URL 下载 Excel 文件，AI 智能识别并预览表结构/数据"
    input_schema = {
        "type": "object",
        "properties": {
            "connectionId": {"type": "string", "description": "数据源 ID"},
            "target": {"type": "string", "enum": ["prod", "test"], "default": "test", "description": "环境，默认 test"},
            "file_url": {"type": "string", "description": "Excel 文件下载地址（.xlsx/.xls, <500KB）"},
        },
        "required": ["connectionId", "file_url"],
    }

    async def execute(self, args: dict) -> dict:
        args = dict(args)
        connection_id = args.pop("connectionId")
        target = args.pop("target", "test")
        file_url = args.pop("file_url")

        # 0. 基础 URL 校验
        parsed_url = urlparse(file_url)
        if parsed_url.scheme not in ("http", "https"):
            raise ValueError("仅支持 http/https 协议的下载地址")

        # 1. 下载文件
        download_resp = await self.client.client.get(file_url, follow_redirects=True, timeout=30.0)
        download_resp.raise_for_status()

        # 1.1 预先根据 Content-Length 拦截超大文件，避免无意义下载
        max_size = 500 * 1024
        content_length = download_resp.headers.get("content-length")
        if content_length:
            try:
                if int(content_length) > max_size:
                    raise ValueError(
                        f"文件 Content-Length {content_length} bytes 超过平台限制 {max_size} bytes (500KB)"
                    )
            except ValueError as e:
                if "超过平台限制" in str(e):
                    raise
                # 非数字时忽略

        file_content = download_resp.content

        # 2. 文件大小校验（平台限制 500KB）
        if len(file_content) > max_size:
            raise ValueError(
                f"文件大小 {len(file_content)} bytes 超过平台限制 {max_size} bytes (500KB)，"
                "请先压缩或拆分后重试"
            )

        # 3. 解析文件名并校验扩展名
        file_name = self._extract_file_name(download_resp, file_url)
        allowed_exts = (".xlsx", ".xls")
        if not file_name.lower().endswith(allowed_exts):
            raise ValueError(
                f"不支持的文件类型: {file_name}，仅支持 {', '.join(allowed_exts)}"
            )

        # 4. 推断 MIME 类型
        content_type, _ = mimetypes.guess_type(file_name)
        if not content_type:
            content_type = (
                "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
                if file_name.lower().endswith(".xlsx")
                else "application/vnd.ms-excel"
            )

        # 5. 构建文件上传
        files = {
            "file": (file_name, io.BytesIO(file_content), content_type),
        }

        return await self.client.upload(
            f"/datasource/connection/{connection_id}/import_document/preview",
            files=files,
            params={"target": target},
        )

    @staticmethod
    def _extract_file_name(response, file_url: str) -> str:
        """从响应头或 URL 中提取文件名。"""
        # 优先从 Content-Disposition 解析
        content_disposition = response.headers.get("content-disposition")
        if content_disposition:
            # 简单解析 filename="xxx" 或 filename*=UTF-8''xxx
            # 注意：filename* 也包含 filename= 前缀，必须先检查 filename*
            for part in content_disposition.split(";"):
                part = part.strip()
                if part.lower().startswith("filename*="):
                    encoded = part.split("=", 1)[1]
                    # 形如 UTF-8''xxx
                    if "''" in encoded:
                        _, _, file_name = encoded.partition("''")
                        return os.path.basename(unquote(file_name, encoding="utf-8"))
                    return os.path.basename(unquote(encoded, encoding="utf-8"))
                if part.lower().startswith("filename="):
                    file_name = part.split("=", 1)[1].strip('"')
                    return os.path.basename(unquote(file_name))

        #  fallback 从 URL 路径获取
        parsed = urlparse(file_url)
        path = unquote(parsed.path)
        file_name = os.path.basename(path) if path else ""
        if file_name and "." in file_name:
            return file_name

        return "import.xlsx"


@register_tool("confirm_import_data")
class ConfirmImportDataTool(ToolDef):
    name = "confirm_import_data"
    description = (
        "确认导入 AI 识别后的数据（建表+插数据）。"
        "传入 preview_import_data 返回的 data 原文 + databaseName 即可，"
        "工具内部会自动组装成后端要求的 {tableStructure(含databaseName), allData} 结构"
    )
    input_schema = {
        "type": "object",
        "properties": {
            "connectionId": {"type": "string", "description": "数据源连接 ID"},
            "databaseName": {"type": "string", "description": "落库的数据库名（导入目标库）"},
            "target": {"type": "string", "enum": ["prod", "test"], "default": "test", "description": "环境，默认 test"},
            "data": {"type": "object", "description": "preview_import_data 返回的 data（含 tableStructure/allData），或已组装好的最终结构"},
        },
        "required": ["connectionId", "data"],
    }

    @staticmethod
    def _build_body(data: dict, database_name) -> dict:
        """把 preview 返回的 data 组装成 confirm 要求的 body。

        后端真实结构（已真机探测确认）：
          { "tableStructure": <单表对象，含 columns 且需带 databaseName>, "allData": [...] }
        preview 返回的 data.tableStructure 是 {success,message,data:{tables:[...]}} 的包装，
        需取出 data.tables[0] 作为单表对象。
        """
        if not isinstance(data, dict):
            return data

        # 若传入的是 preview 工具的完整返回信封 {code, msg, data:{...}}，先剥一层
        if "tableStructure" not in data and isinstance(data.get("data"), dict):
            inner_env = data["data"]
            if "tableStructure" in inner_env or "allData" in inner_env:
                data = inner_env

        ts = data.get("tableStructure")
        single_table = None
        if isinstance(ts, dict):
            if "columns" in ts:
                # 已是单表对象（调用方自行组装过）
                single_table = dict(ts)
            else:
                # preview 包装：tableStructure.data.tables[0]
                inner = ts.get("data") if isinstance(ts.get("data"), dict) else {}
                tables = inner.get("tables") if isinstance(inner, dict) else None
                if isinstance(tables, list) and tables:
                    single_table = dict(tables[0])

        if single_table is None:
            # 无法识别结构，原样透传（兼容调用方已构造好完整 body 的情况）
            return data

        if database_name and not single_table.get("databaseName"):
            single_table["databaseName"] = database_name

        all_data = data.get("allData")
        if all_data is None:
            all_data = data.get("data") or []
        return {"tableStructure": single_table, "allData": all_data}

    async def execute(self, args: dict) -> dict:
        args = dict(args)
        connection_id = args.pop("connectionId")
        target = args.pop("target", "test")
        database_name = args.pop("databaseName", None)
        data = args.pop("data")

        body = self._build_body(data, database_name)
        return await self.client.post(
            f"/datasource/connection/{connection_id}/import_document/confirm",
            json_data=body,
            params={"target": target},
        )
