"""Typed, side-effect-free expressions used by built-in state updates.

Expressions always read one immutable snapshot. Callers validate first, then
evaluate every rule, and only commit after all values succeed.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Mapping


MAX_EXPRESSION_DEPTH = 8
MAX_UPDATE_RULES = 10
MAX_LLM_PARAMETERS = 10


class StateExpressionError(ValueError):
    pass


def _field(value: Mapping[str, Any], snake: str, camel: str | None = None) -> Any:
    if snake in value:
        return value[snake]
    return value.get(camel or snake)


def _value_type(value: Any) -> str:
    if isinstance(value, bool):
        return "boolean"
    if isinstance(value, (int, float)):
        return "number"
    if isinstance(value, str):
        return "string"
    raise StateExpressionError("状态表达式只支持字符串、数字或布尔值")


def _definition_types(
    definitions: Mapping[str, Mapping[str, Any]] | None,
) -> dict[str, str]:
    result: dict[str, str] = {}
    for name, definition in (definitions or {}).items():
        value_type = (
            (definition or {}).get("type")
            if isinstance(definition, Mapping)
            else getattr(definition, "type", None)
        )
        result[str(name)] = str(value_type or "string")
    return result


@dataclass(frozen=True)
class LlmParameter:
    name: str
    value_type: str
    description: str


def _inspect_expression(
    expression: Mapping[str, Any],
    variable_types: Mapping[str, str],
    *,
    allow_llm: bool,
    depth: int,
    parameters: dict[str, LlmParameter],
) -> str:
    if depth > MAX_EXPRESSION_DEPTH:
        raise StateExpressionError(
            f"状态表达式深度不能超过 {MAX_EXPRESSION_DEPTH}"
        )
    kind = str(expression.get("type") or "")
    if kind == "constant":
        return _value_type(expression.get("value"))
    if kind == "variable":
        name = str(_field(expression, "variable_name", "variableName") or "")
        if name not in variable_types:
            raise StateExpressionError(f"状态表达式引用了未声明变量: {name}")
        return variable_types[name]
    if kind == "llm":
        if not allow_llm:
            raise StateExpressionError("确定性 Update State 节点不能使用 LLM 提取值")
        name = str(_field(expression, "parameter_name", "parameterName") or "")
        value_type = str(_field(expression, "value_type", "valueType") or "string")
        description = str(expression.get("description") or "").strip()
        if not name or value_type not in {"string", "number", "boolean"}:
            raise StateExpressionError("LLM 提取参数名称或类型无效")
        if not description:
            raise StateExpressionError(f"LLM 提取参数 {name} 必须填写说明")
        current = LlmParameter(name, value_type, description)
        previous = parameters.get(name)
        if previous is not None and previous != current:
            raise StateExpressionError(f"LLM 参数 {name} 存在冲突定义")
        parameters[name] = current
        if len(parameters) > MAX_LLM_PARAMETERS:
            raise StateExpressionError(
                f"LLM 提取参数不能超过 {MAX_LLM_PARAMETERS} 个"
            )
        return value_type
    if kind == "if":
        condition = expression.get("condition")
        then = expression.get("then")
        otherwise = _field(expression, "otherwise")
        if not all(isinstance(item, Mapping) for item in (condition, then, otherwise)):
            raise StateExpressionError("if 表达式缺少 condition/then/otherwise")
        condition_type = _inspect_expression(
            condition,
            variable_types,
            allow_llm=allow_llm,
            depth=depth + 1,
            parameters=parameters,
        )
        if condition_type != "boolean":
            raise StateExpressionError("if 条件必须是布尔值")
        then_type = _inspect_expression(
            then,
            variable_types,
            allow_llm=allow_llm,
            depth=depth + 1,
            parameters=parameters,
        )
        otherwise_type = _inspect_expression(
            otherwise,
            variable_types,
            allow_llm=allow_llm,
            depth=depth + 1,
            parameters=parameters,
        )
        if then_type != otherwise_type:
            raise StateExpressionError("if 的两个结果分支类型必须一致")
        return then_type

    operands = expression.get("operands")
    if kind not in {
        "and", "or", "eq", "neq", "gt", "gte", "lt", "lte",
        "add", "subtract", "multiply", "divide",
    } or not isinstance(operands, list) or len(operands) < 2:
        raise StateExpressionError(f"状态表达式运算无效: {kind or 'empty'}")
    if not all(isinstance(item, Mapping) for item in operands):
        raise StateExpressionError(f"{kind} 的操作数必须是表达式")
    operand_types = [
        _inspect_expression(
            item,
            variable_types,
            allow_llm=allow_llm,
            depth=depth + 1,
            parameters=parameters,
        )
        for item in operands
    ]
    if kind in {"and", "or"}:
        if any(item != "boolean" for item in operand_types):
            raise StateExpressionError(f"{kind} 只能组合布尔值")
        return "boolean"
    if kind in {"eq", "neq"}:
        if len(set(operand_types)) != 1:
            raise StateExpressionError(f"{kind} 的操作数类型必须一致")
        return "boolean"
    if kind in {"gt", "gte", "lt", "lte"}:
        if any(item != "number" for item in operand_types):
            raise StateExpressionError(f"{kind} 只能比较数字")
        return "boolean"
    if kind == "add" and len(set(operand_types)) == 1:
        result_type = operand_types[0]
        if result_type in {"number", "string"}:
            return result_type
    if any(item != "number" for item in operand_types):
        raise StateExpressionError(f"{kind} 只能用于数字")
    return "number"


def validate_update_rules(
    rules: list[Mapping[str, Any]] | None,
    definitions: Mapping[str, Mapping[str, Any]] | None,
    *,
    allow_llm: bool,
) -> dict[str, dict[str, str]]:
    """Validate rules and return the LLM-visible JSON schema properties."""

    rules = rules or []
    if len(rules) > MAX_UPDATE_RULES:
        raise StateExpressionError(
            f"状态更新规则不能超过 {MAX_UPDATE_RULES} 条"
        )
    variable_types = _definition_types(definitions)
    parameters: dict[str, LlmParameter] = {}
    targets: set[str] = set()
    llm_target_parameters: dict[str, str] = {}
    for rule in rules:
        target = str(_field(rule, "variable_name", "variableName") or "")
        expression = rule.get("expression")
        if target not in variable_types:
            raise StateExpressionError(f"状态更新目标未声明: {target}")
        if target in targets:
            raise StateExpressionError(f"状态更新目标不能重复: {target}")
        targets.add(target)
        if not isinstance(expression, Mapping):
            raise StateExpressionError(f"状态更新 {target} 缺少表达式")
        result_type = _inspect_expression(
            expression,
            variable_types,
            allow_llm=allow_llm,
            depth=1,
            parameters=parameters,
        )
        if result_type != variable_types[target]:
            raise StateExpressionError(
                f"状态更新 {target} 的结果类型应为 {variable_types[target]}"
            )
        if allow_llm:
            llm_names: set[str] = set()
            _collect_llm_names(expression, llm_names)
            for parameter_name in llm_names:
                previous_target = llm_target_parameters.get(parameter_name)
                if previous_target is not None and previous_target != target:
                    raise StateExpressionError(
                        f"LLM 参数 {parameter_name} 不能同时更新多个变量"
                    )
                llm_target_parameters[parameter_name] = target
    return {
        parameter.name: {
            "type": parameter.value_type,
            "description": parameter.description,
        }
        for parameter in parameters.values()
    }


def _collect_llm_names(expression: Mapping[str, Any], result: set[str]) -> None:
    kind = str(expression.get("type") or "")
    if kind == "llm":
        result.add(
            str(_field(expression, "parameter_name", "parameterName") or "")
        )
        return
    if kind == "if":
        children = [
            expression.get("condition"),
            expression.get("then"),
            expression.get("otherwise"),
        ]
    else:
        children = expression.get("operands") or []
    for child in children:
        if isinstance(child, Mapping):
            _collect_llm_names(child, result)


def _evaluate(
    expression: Mapping[str, Any],
    snapshot: Mapping[str, Any],
    arguments: Mapping[str, Any],
) -> Any:
    kind = str(expression.get("type") or "")
    if kind == "constant":
        return expression.get("value")
    if kind == "variable":
        name = str(_field(expression, "variable_name", "variableName") or "")
        if name not in snapshot:
            raise StateExpressionError(f"状态变量尚未赋值: {name}")
        return snapshot[name]
    if kind == "llm":
        name = str(_field(expression, "parameter_name", "parameterName") or "")
        if name not in arguments:
            raise StateExpressionError(f"缺少 LLM 提取参数: {name}")
        return arguments[name]
    if kind == "if":
        branch = "then" if bool(
            _evaluate(expression["condition"], snapshot, arguments)
        ) else "otherwise"
        return _evaluate(expression[branch], snapshot, arguments)

    values = [
        _evaluate(item, snapshot, arguments)
        for item in expression.get("operands") or []
    ]
    if kind == "and":
        return all(values)
    if kind == "or":
        return any(values)
    if kind == "eq":
        return all(value == values[0] for value in values[1:])
    if kind == "neq":
        return all(value != values[0] for value in values[1:])
    if kind == "gt":
        return all(left > right for left, right in zip(values, values[1:]))
    if kind == "gte":
        return all(left >= right for left, right in zip(values, values[1:]))
    if kind == "lt":
        return all(left < right for left, right in zip(values, values[1:]))
    if kind == "lte":
        return all(left <= right for left, right in zip(values, values[1:]))
    if kind == "add":
        result = values[0]
        for value in values[1:]:
            result += value
        return result
    if kind == "subtract":
        result = values[0]
        for value in values[1:]:
            result -= value
        return result
    if kind == "multiply":
        result = values[0]
        for value in values[1:]:
            result *= value
        return result
    if kind == "divide":
        result = values[0]
        for value in values[1:]:
            if value == 0:
                raise StateExpressionError("状态表达式不能除以零")
            result /= value
        return result
    raise StateExpressionError(f"不支持的状态表达式: {kind}")


def evaluate_update_rules(
    rules: list[Mapping[str, Any]] | None,
    snapshot: Mapping[str, Any],
    arguments: Mapping[str, Any] | None = None,
) -> dict[str, Any]:
    """Evaluate every rule from the same snapshot without mutating state."""

    return {
        str(_field(rule, "variable_name", "variableName")): _evaluate(
            rule["expression"], snapshot, arguments or {}
        )
        for rule in rules or []
    }