#!/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 random
import re
import string
import uuid
from datetime import datetime, timedelta
from typing import Any, Callable, Optional


# 占位符正则：{__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))