"""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.schemas.actions import AddAction, AgentAction, UpdateAction
from app.schemas.screen import Screen

QUERY_KEY = "$query"


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

    index: int
    node_id: str
    path: str
    columns: list[str]
    rows: int


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

    add 用了 $query 又没写 sample_data 时设为 false：数据来自用户上传的文件，不是示例数据
    """
    filled: list[AgentAction] = []
    embedded: list[EmbeddedQuery] = []
    for index, action in enumerate(actions):
        found: list[tuple[str, list[str], int]] = []
        node_id = ""  # 只有 add 和 update 有字段值，其他动作不会有替换
        try:
            if isinstance(action, AddAction):
                node = _fill(
                    action.node.model_dump(mode="json", exclude_unset=True),
                    data,
                    "",
                    found,
                )
                if found:
                    fields = action.model_dump(mode="json", exclude_unset=True)
                    fields["node"] = node
                    fields.setdefault("sample_data", False)
                    action = AddAction.model_validate(fields)
                node_id = action.node.id
            elif isinstance(action, UpdateAction):
                value = _fill(action.value, data, action.field_path, found)
                if found:
                    action = action.model_copy(update={"value": value})
                node_id = action.target_id
        except ValueError as error:  # 含 QueryError
            raise BatchError(index, action, str(error)) from error
        embedded.extend(
            EmbeddedQuery(index, node_id, path, columns, rows)
            for path, columns, rows in found
        )
        filled.append(action)
    return filled, embedded


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


def _query(
    value: dict[str, Any],
    data: DataContext,
    path: str,
    found: list[tuple[str, list[str], int]],
) -> list[dict[str, Any]]:
    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 起不同的别名"
        )
    found.append((path, result.columns, len(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 encode_problems(screen: Screen, node_ids: set[str]) -> list[str]:
    """用了 $query 的组件，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"series.{position}.encode.{key} 用了 {'、'.join(missing)}"
                    )
        if wrong:
            problems.append(
                f"组件 {node.id}：{'；'.join(wrong)}，但 dataset.source 的列是 {'、'.join(columns)}；"
                "把 encode 改成查询结果的列名，或在 SQL 里用 AS 起对应的别名"
            )
    return problems