#!/usr/bin/env python3
"""
脚本功能：读取 llm_idea.md 文件，
检查 **Implementation Example** 中的数据字段是否缺少大括号，
如果缺少则添加大括号，然后回写文件。
"""

import re
from pathlib import Path


def get_function_names():
    """获取已知的函数名列表"""
    return {
        'divide', 'subtract', 'add', 'abs', 'power',
        'greater', 'greater_equal', 'equal', 'and', 'or', 'not',
        'ts_delay', 'ts_mean', 'ts_std_dev', 'ts_sum', 'ts_backfill',
        'group_rank', 'group_neutralize', 'zscore',
        'days_from_last_change'
    }


def get_keywords():
    """获取关键字列表"""
    return {
        'True', 'False', 'None', 'and', 'or', 'not', 'if', 'else',
        'for', 'while', 'return', 'in', 'is', 'lambda'
    }


def add_braces_to_expression(expression):
    """
    给表达式中的数据字段添加大括号
    
    使用更健壮的方法：
    1. 找出所有已经被大括号包围的部分，标记为已处理
    2. 找出所有未被包围的变量，添加大括号
    
    Args:
        expression: 原始表达式字符串
        
    Returns:
        str: 添加了大括号后的表达式
    """
    func_names = get_function_names()
    keywords = get_keywords()
    
    # 创建结果列表（用于构建最终字符串）
    result = []
    i = 0
    length = len(expression)
    
    while i < length:
        # 如果当前字符是 '{'，说明已经是大括号包围的内容，直接跳过到匹配的 '}'
        if expression[i] == '{':
            # 找到对应的 '}'
            j = expression.find('}', i)
            if j != -1:
                # 保留原有的大括号内容
                result.append(expression[i:j+1])
                i = j + 1
            else:
                # 没有找到闭合的 '}'，当作普通字符处理
                result.append(expression[i])
                i += 1
        elif expression[i].isalpha() or expression[i] == '_':
            # 可能是变量名的开始
            j = i
            while j < length and (expression[j].isalnum() or expression[j] == '_'):
                j += 1
            
            word = expression[i:j]
            
            # 检查是否是变量（不是函数名、关键字、数字）
            if (word and 
                not word[0].isdigit() and 
                word not in func_names and 
                word not in keywords):
                # 这是一个变量，添加大括号
                result.append('{' + word + '}')
            else:
                # 不是变量，保持原样
                result.append(word)
            
            i = j
        else:
            # 其他字符（数字、运算符、括号等），直接保留
            result.append(expression[i])
            i += 1
    
    return ''.join(result)


def process_file(file_path):
    """
    处理文件，给 Implementation Example 中的数据字段添加大括号
    
    Args:
        file_path: Path对象，指向要处理的文件
        
    Returns:
        bool: 是否成功处理
    """
    try:
        content = file_path.read_text(encoding='utf-8')
    except FileNotFoundError:
        print(f"错误：找不到文件 {file_path}")
        return False
    except Exception as e:
        print(f"读取文件时出错：{e}")
        return False
    
    # 正则匹配 **Implementation Example**: 后面的内容
    pattern = r'(\*\*Implementation Example\*\*:\s*`(.*?)`)'
    
    def replace_match(match):
        """替换函数：对匹配到的内容进行处理"""
        prefix = match.group(1).split('`')[0] + '`'  # **Implementation Example**: `
        original_expr = match.group(2)  # 表达式内容
        suffix = '`'  # 结束的 `
        
        # 给表达式添加大括号
        new_expr = add_braces_to_expression(original_expr)
        
        return prefix + new_expr + suffix
    
    # 执行替换
    new_content = re.sub(pattern, replace_match, content, flags=re.DOTALL)
    
    # 检查是否有变化
    if new_content != content:
        try:
            file_path.write_text(new_content, encoding='utf-8')
            print(f"✓ 成功更新文件：{file_path}")
            return True
        except Exception as e:
            print(f"写入文件时出错：{e}")
            return False
    else:
        print("✓ 文件无需更改（所有数据字段已有大括号）")
        return True


def main():
    # 获取当前目录下的 llm_idea.md 文件
    current_dir = Path(__file__).parent
    file_path = current_dir / "llm_idea.md"
    
    print("=" * 80)
    print("Implementation Example 字段大括号补全工具")
    print("=" * 80)
    print(f"\n正在处理文件：{file_path}\n")
    
    # 处理文件
    success = process_file(file_path)
    
    if success:
        print("\n" + "=" * 80)
        print("处理完成！")
        print("=" * 80)
    else:
        print("\n处理失败！")


if __name__ == "__main__":
    main()