#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
模块名称：template_resolver.py
模块描述：性能测试请求体模板变量替换器

支持占位符：
- {__NOW__}             → 当前时间 "YYYY-MM-DD HH:mm:ss"
- {__NOW_+N_H__}        → 当前时间 + N 小时
- {__NOW_+N_M__}        → 当前时间 + N 分钟
- {__NOW_+N_D__}        → 当前时间 + N 天
- {__NOW_+N_H_INDEX__}  → 当前时间 + N 小时 + 请求序号分钟（并发错峰）
- {__NOW_INDEX__}       → 当前时间 + 请求序号分钟（并发错峰）
- {__RANDOM_NAME_N__}   → N 位随机字母数字
- {__UUID__}            → 随机 UUID（无横线）
- {__RANDOM_INT_N_M__}  → N~M 随机整数
- {__INDEX__}           → 请求序号
- {__REF:task_id:key__} → 引用其他任务输出

作者：czj
创建日期：2026-08-13
"""

import copy
import random
import re
import string
import uuid
from datetime import datetime, timedelta
from typing import Any, Callable, Optional, List, Dict


# 占位符正则：{__TOKEN__}
# token 支持字母、下划线、数字、冒号、加号（如 NOW_+2_H / REF:task:key）
PATTERN = re.compile(r"\{__([A-Za-z_+0-9:]+?)__\}")


class TemplateResolver:
    """
    请求体模板变量替换器

    不对原模板做修改，resolve() 返回新的容器（dict/list）。

    Usage:
        >>> resolver = TemplateResolver()
        >>> body = resolver.resolve({"name": "{__RANDOM_NAME_8__}", "time": "{__NOW__}"})
    """

    def __init__(self, ref_provider: Optional[Callable[[str, str], list]] = None):
        """
        Args:
            ref_provider: 跨任务引用取值回调 (task_id, key) -> list
                          None 时 REF 占位符不替换（返回空字符串）
        """
        self._ref_provider = ref_provider

    def resolve(self, template: Any, request_index: int = 0) -> Any:
        """
        递归替换模板中的占位符

        Args:
            template: dict/list/str，支持嵌套
            request_index: 当前请求序号（对应 {__INDEX__}）

        Returns:
            替换后的 dict/list/str（不修改原模板）
        """
        if isinstance(template, dict):
            return {
                k: self.resolve(v, request_index)
                for k, v in template.items()
            }
        if isinstance(template, list):
            return [self.resolve(v, request_index) for v in template]
        if isinstance(template, str):
            return self._replace_string(template, request_index)
        return template

    def _replace_string(self, text: str, request_index: int) -> str:
        """替换单个字符串中的所有占位符"""
        if "__" not in text:
            return text

        def _repl(match):
            token = match.group(1)
            return self._resolve_token(token, request_index)

        return PATTERN.sub(_repl, text)

    def _resolve_token(self, token: str, request_index: int) -> str:
        """
        解析单个占位符 token

        支持格式：
        - NOW / NOW_+2_H / NOW_+30_M / NOW_+1_D（_ 分隔）
        - NOW_+2_H_INDEX / NOW_INDEX（_INDEX 后缀，追加 request_index 分钟）
        - RANDOM_NAME_6（_ 分隔）
        - RANDOM_INT_1_100（_ 分隔）
        - UUID / INDEX
        - REF:task_id:key（: 分隔）

        Args:
            token: 占位符内部内容（不含 {__ __} 外壳）
            request_index: 当前请求序号

        Returns:
            str: 替换后的值；无法解析时保留原样
        """
        # REF 引用：冒号分隔（REF:task_id:key）
        if token.startswith("REF:"):
            parts = token.split(":")
            if len(parts) >= 3:
                task_id, key = parts[1], parts[2]
                if self._ref_provider:
                    values = self._ref_provider(task_id, key)
                    if values:
                        # 轮询取用，避免并发下都取第一个
                        return str(values[request_index % len(values)])
                # 引用为空或没有 provider → 空串（调用方记录警告）
                return ""
            return f"{{__{token}__}}"

        # 其余占位符按下划线分隔解析
        parts = token.split("_")
        kind = parts[0]

        if kind == "NOW":
            # NOW / NOW_+2_H / NOW_+30_M / NOW_+1_D / 各变体 + _INDEX 后缀
            if len(parts) > 1:
                if parts[1] == "INDEX":
                    # {__NOW_INDEX__} → 当前时间 + request_index 分钟
                    return self._now_with_index_offset("0", "M", request_index)
                delta = parts[1]
                unit = parts[2] if len(parts) > 2 and parts[2] in ("H", "M", "D") else "M"
                if len(parts) > 3 and parts[3] == "INDEX":
                    # {__NOW_+N_H_INDEX__} → 当前时间 + N + request_index 分钟
                    return self._now_with_index_offset(delta, unit, request_index)
                return self._now_with_offset(delta, unit)
            return self._now_str()

        if token.startswith("RANDOM_NAME_"):
            length = int(parts[2]) if len(parts) > 2 else 8
            return self._random_string(length)

        if token.startswith("RANDOM_INT_"):
            try:
                lo, hi = int(parts[2]), int(parts[3])
                if lo > hi:
                    lo, hi = hi, lo
                return str(random.randint(lo, hi))
            except (ValueError, IndexError):
                return f"{{__{token}__}}"

        if kind == "UUID":
            return uuid.uuid4().hex

        if kind == "INDEX":
            return str(request_index)

        # 未知占位符 → 保留原样
        return f"{{__{token}__}}"

    # ---------- 时间相关 ----------

    def _now_str(self) -> str:
        """当前时间字符串 YYYY-MM-DD HH:mm:ss"""
        return datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    def _now_with_offset(self, delta: str, unit: str) -> str:
        """
        当前时间 + 偏移

        Args:
            delta: 偏移量（如 "+2"）
            unit: 单位（H=小时, M=分钟, D=天）

        Returns:
            str: 偏移后的时间字符串
        """
        try:
            n = int(delta.lstrip("+"))
        except ValueError:
            return self._now_str()

        base = datetime.now()
        if unit == "H":
            base += timedelta(hours=n)
        elif unit == "D":
            base += timedelta(days=n)
        else:  # 默认分钟
            base += timedelta(minutes=n)
        return base.strftime("%Y-%m-%d %H:%M:%S")

    def _now_with_index_offset(self, delta: str, unit: str, request_index: int) -> str:
        """
        当前时间 + 偏移 + request_index 分钟（INDEX 错峰）

        复用 _now_with_offset 保证基础偏移逻辑单一来源，
        再追加 request_index 分钟使每个请求时间槽错开。

        Args:
            delta: 偏移量（如 "+2"）
            unit: 单位（H=小时, M=分钟, D=天）
            request_index: 请求序号（追加的分钟数）

        Returns:
            str: 偏移后的时间字符串
        """
        base = self._now_with_offset(delta, unit)
        base_dt = datetime.strptime(base, "%Y-%m-%d %H:%M:%S") + timedelta(minutes=request_index)
        return base_dt.strftime("%Y-%m-%d %H:%M:%S")

    # ---------- 随机值相关 ----------

    @staticmethod
    def _random_string(length: int) -> str:
        """生成 N 位随机字母数字字符串"""
        length = max(length, 1)
        chars = string.ascii_letters + string.digits
        return "".join(random.choice(chars) for _ in range(length))

    # ==================== 唯一性字段 ====================

    @staticmethod
    def apply_unique_fields(body: Any, unique_fields: List[Dict], request_index: int) -> Any:
        """
        对已解析的 body 应用唯一性字段配置。

        在 bodyTemplate 解析完成后调用，对指定字段按请求序号生成唯一值。

        Args:
            body: 已解析的请求体（dict/list/str）
            unique_fields: 唯一性字段配置列表，每项含 field/strategy/template
            request_index: 当前请求序号（从 0 递增）

        Returns:
            唯一化后的 body（不修改原对象）
        """
        if not unique_fields:
            return body

        # 深拷贝 body
        result = copy.deepcopy(body)

        for rule in unique_fields:
            field_path = rule.get("field", "")
            strategy = rule.get("strategy", "suffix")
            template = rule.get("template")

            value = _get_field_by_path(result, field_path)
            if value is not None and isinstance(value, str):
                new_value = _generate_unique_value(value, strategy, template, request_index)
                _set_field_by_path(result, field_path, new_value)

        return result


def _get_field_by_path(obj: Any, path: str) -> Any:
    """
    通过 JSON Path 获取字段值。

    支持格式：
    - 'templateName' → 简单字段名
    - '$[0].templateName' → 顶层数组第 0 个元素的字段
    - 'topicList[0].topicName' → 嵌套数组字段
    - 'level1.level2' → 深层嵌套字段

    对数组会透明遍历：'topicList[0].topicName' 作用于顶层 list 时，
    会依次应用到每个元素并返回第一个命中值。
    """
    clean_path = re.sub(r"^\$\.?|^\.", "", path)
    if not clean_path:
        return obj

    parts = clean_path.split(".")
    results = [obj]

    for part in parts:
        bare = re.match(r"^\[(\d+)\]$", part)          # 裸下标：[0]
        keyed = re.match(r"^(\w+)\[(\d+)\]$", part)    # 键+下标：topicList[0]

        next_results = []
        for node in results:
            if node is None:
                continue
            if bare:
                idx = int(bare.group(1))
                if isinstance(node, (list, tuple)) and idx < len(node):
                    next_results.append(node[idx])
            elif keyed:
                key, idx = keyed.group(1), int(keyed.group(2))
                if isinstance(node, dict) and key in node:
                    sub = node[key]
                    if isinstance(sub, (list, tuple)) and idx < len(sub):
                        next_results.append(sub[idx])
                elif isinstance(node, (list, tuple)):
                    # 顶层是数组：对每个元素递归查找该部分
                    for item in node:
                        v = _get_field_by_path(item, part)
                        if v is not None:
                            next_results.append(v)
            else:
                if isinstance(node, dict):
                    v = node.get(part)
                    if v is not None:
                        next_results.append(v)
                elif isinstance(node, (list, tuple)):
                    # 当前是数组：对每个元素递归查找该部分
                    for item in node:
                        v = _get_field_by_path(item, part)
                        if v is not None:
                            next_results.append(v)
        results = next_results
        if not results:
            break

    return results[0] if results else None


def _set_field_by_path(obj: Any, path: str, value: Any) -> None:
    """
    通过 JSON Path 设置字段值（原地修改）。

    与 _get_field_by_path 相同的路径语法；路径不存在时静默跳过。
    """
    clean_path = re.sub(r"^\$\.?|^\.", "", path)
    if not clean_path:
        return

    parts = clean_path.split(".")

    def _apply(node: Any, idx: int) -> None:
        """递归：node 是 parts[idx] 对应的容器，idx 是当前要处理的段"""
        part = parts[idx]

        if idx == len(parts) - 1:
            # 最后一段：实际赋值
            bare = re.match(r"^\[(\d+)\]$", part)
            keyed = re.match(r"^(\w+)\[(\d+)\]$", part)
            if bare:
                i = int(bare.group(1))
                if isinstance(node, (list, tuple)) and i < len(node):
                    node[i] = value
            elif keyed:
                key, i = keyed.group(1), int(keyed.group(2))
                if isinstance(node, dict) and key in node:
                    sub = node[key]
                    if isinstance(sub, (list, tuple)) and i < len(sub):
                        sub[i] = value
            elif isinstance(node, dict):
                node[part] = value
            elif isinstance(node, (list, tuple)):
                for item in node:
                    _apply(item, idx)
            return

        bare = re.match(r"^\[(\d+)\]$", part)
        keyed = re.match(r"^(\w+)\[(\d+)\]$", part)
        if bare:
            i = int(bare.group(1))
            if isinstance(node, (list, tuple)) and i < len(node):
                _apply(node[i], idx + 1)
        elif keyed:
            key, i = keyed.group(1), int(keyed.group(2))
            if isinstance(node, dict) and key in node:
                _apply(node[key], idx + 1)
            elif isinstance(node, (list, tuple)):
                for item in node:
                    _apply(item, idx)
        else:
            if isinstance(node, dict):
                if part in node:
                    _apply(node[part], idx + 1)
            elif isinstance(node, (list, tuple)):
                for item in node:
                    _apply(item, idx)

    _apply(obj, 0)


def _generate_unique_value(original: str, strategy: str, template: Optional[str], index: int) -> str:
    """根据策略生成唯一值。"""
    if strategy == "suffix":
        # 追加 _req_{index}（3位补零）
        return f"{original}_req_{index:03d}"
    elif strategy == "template" and template:
        # 使用模板替换（如 "{__UUID__}" → 替换原值）
        return template  # 占位符由后续的 resolve() 处理
    return original