"""
Excel 表格分析：下载、schema 探测、pandas 沙箱执行。

设计原则：
  - MCP 容器内执行，与 backend 进程隔离（建议 Docker 部署）
  - 仅允许对已加载的 ``df`` 做 pandas 运算，禁止 import / 文件 / 网络
  - 代码必须将最终答案赋给 ``result``
"""
from __future__ import annotations

import ast
import json
import tempfile
from pathlib import Path
from typing import Any

import httpx
import pandas as pd

MAX_DOWNLOAD_BYTES = 20 * 1024 * 1024
DOWNLOAD_TIMEOUT = 120.0
EXEC_TIMEOUT_SEC = 15
MAX_RESULT_ROWS = 200

FORBIDDEN_CALL_NAMES = frozenset(
    {
        "open",
        "exec",
        "eval",
        "compile",
        "__import__",
        "getattr",
        "setattr",
        "delattr",
        "globals",
        "locals",
        "vars",
        "dir",
        "help",
        "input",
        "breakpoint",
        "exit",
        "quit",
    }
)

FORBIDDEN_ATTR_NAMES = frozenset({"__import__", "__subclasses__", "__globals__", "__code__"})

SAFE_BUILTINS = {
    "len": len,
    "str": str,
    "int": int,
    "float": float,
    "list": list,
    "dict": dict,
    "tuple": tuple,
    "set": set,
    "sum": sum,
    "min": min,
    "max": max,
    "range": range,
    "round": round,
    "abs": abs,
    "sorted": sorted,
    "enumerate": enumerate,
    "zip": zip,
    "map": map,
    "filter": filter,
    "any": any,
    "all": all,
    "bool": bool,
    "True": True,
    "False": False,
    "None": None,
}


def _validate_pandas_code(code: str) -> None:
    if not code or not code.strip():
        raise ValueError("pandas_code 不能为空")
    if len(code) > 8000:
        raise ValueError("pandas_code 过长（上限 8000 字符）")

    tree = ast.parse(code, mode="exec")

    for node in ast.walk(tree):
        if isinstance(node, (ast.Import, ast.ImportFrom)):
            raise ValueError("禁止 import，请仅使用已提供的 pd 与 df")
        if isinstance(node, ast.Call):
            func = node.func
            if isinstance(func, ast.Name) and func.id in FORBIDDEN_CALL_NAMES:
                raise ValueError(f"禁止调用: {func.id}")
            if isinstance(func, ast.Attribute) and func.attr in FORBIDDEN_ATTR_NAMES:
                raise ValueError(f"禁止访问: {func.attr}")
        if isinstance(node, ast.Attribute) and node.attr in FORBIDDEN_ATTR_NAMES:
            raise ValueError(f"禁止访问属性: {node.attr}")


def _download_excel(file_url: str) -> Path:
    url = (file_url or "").strip()
    if not url.startswith(("http://", "https://")):
        raise ValueError("file_url 必须是 http(s) 可下载地址")

    suffix = ".xlsx"
    lower = url.split("?", 1)[0].lower()
    if lower.endswith(".xls"):
        suffix = ".xls"
    elif lower.endswith(".csv"):
        suffix = ".csv"

    with httpx.Client(timeout=DOWNLOAD_TIMEOUT, follow_redirects=True) as client:
        with client.stream("GET", url) as resp:
            resp.raise_for_status()
            cl = resp.headers.get("content-length")
            if cl and int(cl) > MAX_DOWNLOAD_BYTES:
                raise ValueError(f"文件过大（>{MAX_DOWNLOAD_BYTES // (1024 * 1024)}MB）")

            tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
            try:
                total = 0
                for chunk in resp.iter_bytes():
                    total += len(chunk)
                    if total > MAX_DOWNLOAD_BYTES:
                        raise ValueError(f"文件过大（>{MAX_DOWNLOAD_BYTES // (1024 * 1024)}MB）")
                    tmp.write(chunk)
                tmp.flush()
                return Path(tmp.name)
            except Exception:
                tmp.close()
                Path(tmp.name).unlink(missing_ok=True)
                raise
            finally:
                tmp.close()


def _resolve_sheet_name(sheet_name: str, excel_path: Path) -> str | int:
    raw = (sheet_name or "").strip()
    xl = pd.ExcelFile(excel_path)
    if not raw or raw.lower() in ("0", "first", "default"):
        return xl.sheet_names[0] if xl.sheet_names else 0
    if raw.isdigit():
        idx = int(raw)
        if idx < 0 or idx >= len(xl.sheet_names):
            raise ValueError(f"sheet 索引越界: {idx}，共 {len(xl.sheet_names)} 个 sheet")
        return xl.sheet_names[idx]
    if raw not in xl.sheet_names:
        raise ValueError(f"sheet 不存在: {raw}，可选: {xl.sheet_names}")
    return raw


def _load_dataframe(excel_path: Path, sheet_name: str) -> pd.DataFrame:
    if excel_path.suffix.lower() == ".csv":
        return pd.read_csv(excel_path)
    resolved = _resolve_sheet_name(sheet_name, excel_path)
    return pd.read_excel(excel_path, sheet_name=resolved)


def _format_result(result: Any) -> str:
    if isinstance(result, pd.DataFrame):
        df = result.head(MAX_RESULT_ROWS)
        payload = {
            "type": "dataframe",
            "rows": len(result),
            "truncated": len(result) > MAX_RESULT_ROWS,
            "data": json.loads(df.to_json(orient="records", force_ascii=False)),
        }
        return json.dumps(payload, ensure_ascii=False)
    if isinstance(result, pd.Series):
        s = result.head(MAX_RESULT_ROWS)
        payload = {
            "type": "series",
            "length": len(result),
            "truncated": len(result) > MAX_RESULT_ROWS,
            "data": json.loads(s.to_json(force_ascii=False)),
        }
        return json.dumps(payload, ensure_ascii=False)
    if isinstance(result, (int, float, bool)) or result is None:
        return json.dumps({"type": "scalar", "value": result}, ensure_ascii=False)
    if isinstance(result, str):
        return json.dumps({"type": "text", "value": result[:5000]}, ensure_ascii=False)
    return json.dumps({"type": "other", "value": str(result)[:5000]}, ensure_ascii=False)


def _sandbox_execute(code: str, df: pd.DataFrame) -> Any:
    _validate_pandas_code(code)
    namespace: dict[str, Any] = {"pd": pd, "df": df, "result": None}
    compiled = compile(code, "<pandas_sandbox>", "exec")

    def _run() -> None:
        exec(compiled, {"__builtins__": SAFE_BUILTINS}, namespace)  # noqa: S102

    import concurrent.futures

    with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
        fut = pool.submit(_run)
        try:
            fut.result(timeout=EXEC_TIMEOUT_SEC)
        except concurrent.futures.TimeoutError as e:
            raise TimeoutError(f"执行超时（>{EXEC_TIMEOUT_SEC}s）") from e

    if namespace.get("result") is None:
        raise ValueError("请在 pandas_code 末尾将答案赋给变量 result，例如: result = df['列'].sum()")
    return namespace["result"]


def get_excel_schema(file_url: str, sheet_name: str = "") -> str:
    """
    下载 Excel/CSV 并返回 sheet 列表、列名、类型与前几行预览（JSON 字符串）。
    """
    path = _download_excel(file_url)
    try:
        if path.suffix.lower() == ".csv":
            df = pd.read_csv(path, nrows=5)
            preview_rows = json.loads(
                df.head(5).to_json(orient="records", force_ascii=False)
            )
            payload = {
                "format": "csv",
                "sheets": ["default"],
                "default_sheet": "default",
                "schemas": {
                    "default": {
                        "columns": [str(c) for c in df.columns],
                        "dtypes": {str(c): str(df[c].dtype) for c in df.columns},
                        "preview_rows": preview_rows,
                    }
                },
            }
            return json.dumps(payload, ensure_ascii=False)

        xl = pd.ExcelFile(path)
        schemas: dict[str, Any] = {}
        for name in xl.sheet_names:
            df = pd.read_excel(path, sheet_name=name, nrows=5)
            schemas[name] = {
                "columns": [str(c) for c in df.columns],
                "dtypes": {str(c): str(df[c].dtype) for c in df.columns},
                "preview_rows": json.loads(
                    df.head(5).to_json(orient="records", force_ascii=False)
                ),
            }
        payload = {
            "format": "excel",
            "sheets": xl.sheet_names,
            "default_sheet": xl.sheet_names[0] if xl.sheet_names else "",
            "schemas": schemas,
            "hint": "分析时先选 sheet_name，再调用 run_pandas_on_excel；代码中只用 df 与 pd，结果赋给 result",
        }
        if sheet_name.strip():
            resolved = _resolve_sheet_name(sheet_name, path)
            key = resolved if isinstance(resolved, str) else xl.sheet_names[int(resolved)]
            payload["selected_sheet"] = key
        return json.dumps(payload, ensure_ascii=False)
    finally:
        path.unlink(missing_ok=True)


def run_pandas_on_excel(file_url: str, pandas_code: str, sheet_name: str = "0") -> str:
    """
    在沙箱中对 Excel/CSV 执行 pandas 代码。

    pandas_code 示例::
        result = df[df['区域'] == '华东']['销售额'].sum()
        result = df.groupby('部门')['金额'].mean().sort_values(ascending=False).head(10)
    """
    path = _download_excel(file_url)
    try:
        df = _load_dataframe(path, sheet_name)
        if len(df) > 100_000:
            raise ValueError("表格行数过多（>10万），请让用户拆分文件或缩小范围")
        result = _sandbox_execute(pandas_code, df)
        return _format_result(result)
    finally:
        path.unlink(missing_ok=True)