"""数据文件的本地工具：describe_data 查看列概况，query_data 执行只读 SQL；问答和 screen agent 共用

apply_actions 里的 $query 也通过 DataContext.query 执行，见 tools.py

工具只能看到 DataContext 里的表，也就是当前用户在本会话里未过期的文件，查询引擎也只为这些表建视图
工具通过 runtime.context.data 读取 DataContext，用到它们的 agent 上下文都要有 data 字段；
上下文类型因 agent 而异，所以写 ToolRuntime[Any]（不写类型参数时按 None 处理，序列化会告警）
"""

from dataclasses import dataclass
from pathlib import Path
from typing import Annotated, Any

from langchain.tools import ToolRuntime, tool
from langchain_core.messages import ToolMessage
from pydantic import Field

from app.config import settings
from app.datafiles.query import MAX_ROWS, QueryError, QueryResult, run_query
from app.schemas.data_file import RequestDataFile, RequestDataTable

# 列超过这么多时，不指定 columns 只返回列名和类型，避免一次塞进太多概况
BRIEF_COLUMNS = 50


@dataclass(frozen=True)
class DataContext:
    files: tuple[RequestDataFile, ...] = ()

    def find(self, name: str) -> tuple[RequestDataFile, RequestDataTable] | None:
        for file in self.files:
            for table in file.tables:
                if table.table == name:
                    return file, table
        return None

    def paths(self) -> dict[str, Path]:
        return {
            table.table: Path(table.path)
            for file in self.files
            for table in file.tables
        }

    def query(self, sql: str, *, max_rows: int = MAX_ROWS) -> QueryResult:
        """执行一条只读 SQL；失败时抛 QueryError，原因可以直接返回给模型修正"""
        paths = self.paths()
        # 任务执行期间被删除或清理的文件跳过，不影响查询其他表
        existing = {name: path for name, path in paths.items() if path.exists()}
        try:
            return run_query(
                sql, existing, timeout=settings.data_query_timeout, max_rows=max_rows
            )
        except QueryError as error:
            if gone := [name for name in paths if name not in existing]:
                raise QueryError(
                    f"{error}；{'、'.join(gone)} 的文件已删除或过期，不能再查询"
                ) from None
            raise


def _error(runtime: ToolRuntime, message: str) -> ToolMessage:
    # 返回给模型修正，不中断 agent
    return ToolMessage(message, tool_call_id=runtime.tool_call_id, status="error")


def _table_list(data: DataContext) -> str:
    return "、".join(data.paths()) or "无"


@tool
def describe_data(
    table: str,
    runtime: ToolRuntime[Any],
    columns: Annotated[
        list[str] | None,
        Field(description="只看这些列；不传时看全部列"),
    ] = None,
) -> dict[str, Any] | ToolMessage:
    """查看一张数据表的列概况（类型、空值率、去重值个数、数值范围和四分位、常见值、样例）、随机样例行，以及解析时做的调整

    分析或取数前先看要用的表；表名见请求 data_files 里的 table。统计值来自全表，样例只是其中几行
    列很多时只返回列名和类型，再用 columns 指定要看的列
    """
    data: DataContext = runtime.context.data
    found = data.find(table)
    if found is None:
        return _error(runtime, f"没有表 {table}，可用的表：{_table_list(data)}")
    file, item = found
    profiled: list[dict[str, Any]] = item.profile["columns"]
    result: dict[str, Any] = {
        "table": item.table,
        "label": item.label,
        "rows": item.rows,
    }
    if columns is None and len(profiled) > BRIEF_COLUMNS:
        result["columns"] = [
            {key: column[key] for key in ("name", "type", "format") if key in column}
            for column in profiled
        ]
        result["note"] = (
            f"共 {len(profiled)} 列，只列出列名和类型；用 columns 指定列查看详细概况和样例"
        )
    else:
        index = {column["name"]: position for position, column in enumerate(profiled)}
        names = list(index) if columns is None else list(dict.fromkeys(columns))
        positions = [index[name] for name in names if name in index]
        result["columns"] = [profiled[position] for position in positions]
        # 样例行的值和 columns 一一对应
        result["sample_rows"] = [
            [row[position] for position in positions] for row in item.profile["sample"]
        ]
        if missing := [name for name in names if name not in index]:
            result["missing_columns"] = missing
    # 针对这张表的警告，以及针对整个文件的警告（例如跳过的坏行）
    result["warnings"] = [
        warning.message
        for warning in file.warnings
        if warning.table in (item.table, None)
    ]
    return result


@tool
def query_data(sql: str, runtime: ToolRuntime[Any]) -> dict[str, Any] | ToolMessage:
    """对上传的数据表执行一条只读 SQL（DuckDB 语法），表名是请求 data_files 里的 table（t1、t2……）

    中文或带符号的列名用双引号括起来；宽表转长表用 UNPIVOT，PIVOT 要用 IN 写明取值
    最多返回 200 行，truncated 为 true 时结果不完整，改用聚合或 LIMIT
    回答和图表里的数字以查询结果为准，不要自己心算
    """
    data: DataContext = runtime.context.data
    try:
        return data.query(sql).to_dict()
    except QueryError as error:
        return _error(runtime, str(error))


DATA_TOOLS = [describe_data, query_data]