"""apply_actions 里的 $query：把字段值写成 {"$query": "SELECT ..."}，提交时替换成上传数据的查询结果

模型不用把查询结果逐行抄进动作，又慢又容易抄错；替换发生在 run_batch 之前，替换后照常做字段校验
结果的形状按写入字段的定义：单个值的字段（指标卡的数值等）取一行一列的那个值，元素是数字、文字的数组
（指标卡的趋势数据）取唯一一列；其余是对象数组，键是查询结果的列名，正好对应图表物料的 dataset.source，
series 的 encode 按列名引用
"""

from dataclasses import dataclass
from typing import Any

from pydantic import JsonValue

from app.actions.batch import BatchError
from app.agent.data_tools import DataContext
from app.config import settings
from app.resources.material_fields import component_fields, find_material
from app.schemas.actions import AddAction, AgentAction, UpdateAction
from app.schemas.component_resource import ComponentCatalog
from app.schemas.material_field import MaterialField
from app.schemas.screen import Screen

QUERY_KEY = "$query"
# 图表物料的数据都在这里，series 的 encode 按它的列名引用；整个换成查询结果时数据就不再是示例数据
SOURCE_PATH = "props.option.dataset.source"
# 字段是这些类型时，$query 的结果取一行一列的值
SCALAR_TYPES = {"string", "number", "integer", "boolean", "color", "enum"}
# 数组元素是这些类型时，$query 的结果取唯一一列
LIST_ITEM_TYPES = {"string", "number", "integer", "boolean", "color"}


@dataclass(frozen=True)
class EmbeddedQuery:
    """一处替换：第几个动作、写到哪个字段、结果的列和行数，写进 apply_actions 的结果供模型核对

    写入单个值的字段时 value 是取到的值，模型据此核对数字、在回复里引用
    """

    index: int
    node_id: str
    path: str
    columns: list[str]
    rows: int
    value: JsonValue = None


Found = list[EmbeddedQuery]


def fill_queries(
    actions: list[AgentAction],
    data: DataContext,
    catalog: ComponentCatalog | None = None,
    screen: Screen | None = None,
) -> tuple[list[AgentAction], list[EmbeddedQuery]]:
    """替换 add 的 node 和 update 的 value 里所有的 $query；查询失败抛 BatchError，整批不生效

    catalog 和 screen 用来查写入字段的定义，决定结果的形状；没有时都按对象数组
    """
    # update 的目标可能是这一批前面新增的组件，类型按执行顺序记下
    types = {node.id: node.type for node in screen.nodes} if screen else {}
    filled: list[AgentAction] = []
    embedded: list[EmbeddedQuery] = []
    for index, action in enumerate(actions):
        found: Found = []
        try:
            if isinstance(action, AddAction):
                types[action.node.id] = action.node.type
                node = _fill(
                    action.node.model_dump(mode="json", exclude_unset=True),
                    _Target(
                        data, _fields(catalog, action.node.type), index, action.node.id
                    ),
                    "",
                    found,
                )
                if found:
                    fields = action.model_dump(mode="json", exclude_unset=True)
                    fields["node"] = node
                    action = AddAction.model_validate(fields)
            elif isinstance(action, UpdateAction):
                target = _Target(
                    data,
                    _fields(catalog, types.get(action.target_id)),
                    index,
                    action.target_id,
                )
                value = _fill(action.value, target, action.field_path, found)
                if found:
                    action = action.model_copy(update={"value": value})
        except ValueError as error:  # 含 QueryError
            raise BatchError(index, action, str(error)) from error
        embedded.extend(found)
        filled.append(action)
    return filled, embedded


@dataclass(frozen=True)
class _Target:
    """一个动作的替换上下文：查询用的数据、组件的字段定义、动作序号和组件 ID"""

    data: DataContext
    fields: dict[str, MaterialField]
    index: int
    node_id: str


def _fields(
    catalog: ComponentCatalog | None, node_type: str | None
) -> dict[str, MaterialField]:
    """组件类型的字段定义，按路径索引；没有目录或类型不认识时为空，结果按对象数组处理"""
    if catalog is None or node_type is None:
        return {}
    try:
        component = find_material(catalog, node_type)
    except ValueError:  # 物料不存在，补全或执行时会报错
        return {}
    return {field.path: field for field in component_fields(catalog, component)}


def _fill(value: JsonValue, target: _Target, path: str, found: Found) -> JsonValue:
    if isinstance(value, dict):
        if QUERY_KEY in value:
            return _query(value, target, path, found)
        return {
            key: _fill(item, target, _join(path, key), found)
            for key, item in value.items()
        }
    if isinstance(value, list):
        return [
            _fill(item, target, _join(path, position), found)
            for position, item in enumerate(value)
        ]
    return value


def _query(
    value: dict[str, Any], target: _Target, path: str, found: Found
) -> JsonValue:
    data = target.data
    sql = value[QUERY_KEY]
    if len(value) != 1 or not isinstance(sql, str):
        raise ValueError(
            f'{path} 的 $query 要写成 {{"$query": "SELECT ..."}}，对象里不能有其他键'
        )
    if not data.files:
        raise ValueError("会话里没有上传的数据文件，不能使用 $query")
    limit = settings.data_embed_max_rows
    result = data.query(sql, max_rows=limit)
    if result.truncated:
        raise ValueError(
            f"{path} 的 $query 结果超过 {limit} 行，组件放不下；在 SQL 里聚合、筛选或 LIMIT"
        )
    if duplicated := sorted(
        {name for name in result.columns if result.columns.count(name) > 1}
    ):
        raise ValueError(
            f"{path} 的 $query 结果有重名的列 {'、'.join(duplicated)}，用 AS 起不同的别名"
        )
    shape = (
        f"{len(result.rows)} 行 {len(result.columns)} 列（{'、'.join(result.columns)}）"
    )
    field = target.fields.get(path)
    if field is not None and field.value_type in SCALAR_TYPES:
        if len(result.rows) != 1 or len(result.columns) != 1:
            raise ValueError(
                f"{path}（{field.label}）是单个值，$query 的结果要一行一列，实际 {shape}"
            )
        scalar = result.rows[0][0]
        found.append(
            EmbeddedQuery(target.index, target.node_id, path, result.columns, 1, scalar)
        )
        return scalar
    listed = (
        field is not None
        and field.value_type == "array"
        and field.item_type in LIST_ITEM_TYPES
    )
    if listed and len(result.columns) != 1:
        raise ValueError(
            f"{path}（{field.label}）是一组值，$query 的结果要只有一列，实际 {shape}；"
            "排序用的列不用选出来，写在 ORDER BY 里"
        )
    found.append(
        EmbeddedQuery(
            target.index, target.node_id, path, result.columns, len(result.rows)
        )
    )
    if listed:
        return [row[0] for row in result.rows]
    return [dict(zip(result.columns, row, strict=True)) for row in result.rows]


def _join(path: str, key: str | int) -> str:
    return f"{path}.{key}" if path else str(key)


def _binding(node) -> tuple[Any, Any]:
    option = node.props.get("option")
    if not isinstance(option, dict):
        return None, None
    return option.get("dataset"), option.get("series")


def rebound_nodes(before: Screen, after: Screen) -> set[str]:
    """这一批改过数据或映射（option.dataset、option.series）的组件，新增的也算

    只按前后状态比较，不看是哪种动作改的：$query、直接写数据、单独改 encode 都会让两者对不上；没动过的旧组件不核对
    """
    old = {node.id: _binding(node) for node in before.nodes}
    return {
        node.id
        for node in after.nodes
        if node.id not in old or old[node.id] != _binding(node)
    }


def encode_problems(screen: Screen, node_ids: set[str]) -> list[str]:
    """指定组件的 series encode 引用的列必须在 dataset.source 里

    数据换了列名、或 encode 改成了不存在的列时图表是空的，前端也不报错，所以在这里拦下
    """
    problems = []
    for node in screen.nodes:
        if node.id not in node_ids:
            continue
        option = node.props.get("option")
        if not isinstance(option, dict):
            continue
        dataset = option.get("dataset")
        source = dataset.get("source") if isinstance(dataset, dict) else None
        # 空结果或不是对象数组时无从核对
        if not (
            isinstance(source, list)
            and source
            and all(isinstance(row, dict) for row in source)
        ):
            continue
        columns = list(dict.fromkeys(key for row in source for key in row))
        series = option.get("series")
        items = series if isinstance(series, list) else [series]
        wrong = []
        for position, item in enumerate(items):
            encode = item.get("encode") if isinstance(item, dict) else None
            if not isinstance(encode, dict):
                continue
            for key, names in encode.items():
                missing = [
                    name
                    for name in (names if isinstance(names, list) else [names])
                    if isinstance(name, str) and name not in columns
                ]
                if missing:
                    wrong.append(
                        f"props.option.series.{position}.encode.{key} 用了 {'、'.join(missing)}"
                    )
        if wrong:
            problems.append(
                f"组件 {node.id}：{'；'.join(wrong)}，但 props.option.dataset.source 的列是 {'、'.join(columns)}；"
                "把 encode 改成 dataset.source 里的列名；数据来自 $query 时也可以在 SQL 里用 AS 起对应的别名"
            )
    return problems