"""NL2SQL 查询结果脱敏。"""
from __future__ import annotations

import hashlib
from typing import Any


def _mask_partial(value: Any) -> Any:
    if value is None:
        return None
    text = str(value)
    if len(text) == 1:
        return "*"
    if len(text) <= 4:
        return f"{text[0]}{'*' * (len(text) - 2)}{text[-1]}"
    return f"{text[:3]}{'*' * max(1, len(text) - 5)}{text[-2:]}"


def _mask_value(value: Any, mask_type: str) -> Any:
    if value is None:
        return None
    if mask_type == "hash":
        return hashlib.sha256(str(value).encode("utf-8")).hexdigest()
    return _mask_partial(value)


def mask_rows(
    rows: list[dict[str, Any]],
    masks: dict[tuple[str, str], str],
    *,
    table_name: str | None = None,
) -> list[dict[str, Any]]:
    """复制并脱敏结果行，不修改数据库查询返回的原始对象。"""
    output = [dict(row) for row in rows]
    for (table, column), mask_type in masks.items():
        if table_name is not None and table != table_name:
            continue
        for row in output:
            if column in row:
                row[column] = _mask_value(row[column], mask_type)
    return output