import ast
import tokenize
import io
import re


def remove_comments_and_docstrings(source):
    """
    删除 Python 源码中的注释和文档字符串。
    """
    src = io.StringIO(source)
    out = []
    prev_tok_type = tokenize.INDENT
    last_lineno = -1
    last_col = 0

    for tok in tokenize.generate_tokens(src.readline):
        token_type = tok.type
        token_string = tok.string
        start_line, start_col = tok.start
        end_line, end_col = tok.end

        if start_line > last_lineno:
            out.append("\n" * (start_line - last_lineno - 1))
            last_col = 0
        elif start_col > last_col:
            out.append(" " * (start_col - last_col))

        if token_type == tokenize.COMMENT:
            pass
        elif token_type == tokenize.STRING:
            if prev_tok_type not in (tokenize.INDENT, tokenize.NEWLINE):
                # 判断是否为 docstring
                if re.match(r'^\s*"""(?:[^"]|"{1,2})*"""$', token_string) or re.match(
                        r"^\s*'''(?:[^']|'{1,2})*'''$", token_string):
                    continue
                else:
                    out.append(token_string)
            else:
                continue
        else:
            out.append(token_string)

        prev_tok_type = token_type
        last_col = end_col
        last_lineno = end_line

    return "".join([i for i in out if i.strip() != ""]).strip()


def get_method_source(filename, method_name):
    """
    提取指定类中的方法源码。
    """
    with open(filename, "r", encoding="utf-8") as f:
        source = f.read()
    tree = ast.parse(source)

    for node in ast.walk(tree):
        if isinstance(node, ast.ClassDef):
            # 遍历类中的所有方法
            for item in node.body:
                if isinstance(item, ast.FunctionDef) and item.name == method_name:
                    lines = source.splitlines(keepends=True)
                    start_line = item.lineno - 1
                    end_line = _get_function_end_line(item, lines)
                    method_source = "".join(lines[start_line:end_line])
                    return method_source

    raise ValueError(f"Method '{method_name}' not found.")


def _get_function_end_line(node, lines):
    last_child = None
    for child in ast.walk(node):
        if hasattr(child, "lineno"):
            if last_child is None or child.lineno > last_child.lineno:
                last_child = child
    if last_child:
        return last_child.lineno
    return node.lineno


def save_to_tmp(method_source, tmp_file="tmp.txt"):
    with open(tmp_file, "w", encoding="utf-8") as f:
        f.write(method_source)


def read_from_new(new_file="new.txt"):
    with open(new_file, "r", encoding="utf-8") as f:
        return f.read()


def normalize_code(code):
    code_no_comments = remove_comments_and_docstrings(code)
    return code_no_comments


def compare_functions(code1, code2):
    normalized1 = normalize_code(code1)
    normalized2 = normalize_code(code2)
    return normalized1 == normalized2


def replace_method_in_file_with_comments(filename, method_name, new_method_source):
    """
    将指定类中的方法替换为新内容（保留注释等原始结构）。
    """
    with open(filename, "r", encoding="utf-8") as f:
        lines = f.readlines()

    with open(filename, "r", encoding="utf-8") as f:
        source = f.read()

    tree = ast.parse(source)

    for node in ast.walk(tree):
        if isinstance(node, ast.ClassDef):
            for item in node.body:
                if isinstance(item, ast.FunctionDef) and item.name == method_name:
                    start_line = item.lineno - 1  # lineno 是 1-based
                    end_line = _get_function_end_line(item, lines)

                    # 新方法内容按行分割，并保留缩进结构
                    # 注意：new_method_source 应当是 new.txt 的原始字符串
                    new_lines = new_method_source.splitlines(keepends=True)

                    # 替换对应行区间
                    lines[start_line:end_line] = new_lines

                    # 写回文件
                    with open(filename, "w", encoding="utf-8") as f:
                        f.writelines(lines)
                    return

    raise ValueError(f"Method '{method_name}' not found.")


system_prompt = """
角色：你是一个专业的程序员，具有深厚的 Python 编程能力，能够快速理解和修改代码。你会根据我的要求来为代码添加相应的解释。
任务：我会给你一段 Python 代码，你需要根据我的要求添加注释。注意，你不能对代码内容进行任何修改，只能添加注释。
- 在 play_once 中：
    - 你需要理解下面这些可用的函数调用：
        - 我定义了 self.move() 函数，可以传递至多两个参数，分别是 actions_by_arm1, actions_by_arm2。
        - 以下函数会返回 actions（即一个动作序列）：
            - self.grasp_actor(actor, arm_tag:ArmTag, **args)
            - self.place_actor(actor, arm_tag:ArmTag, target_pose, **args)
            - self.move_to_pose(arm_tag:ArmTag, target_pose)
            - self.move_by_displacement(arm_tag:ArmTag, x, y, z)
            - self.close_gripper(arm_tag:ArmTag, **args)
            - self.open_gripper(arm_tag:ArmTag, **args)
            - self.back_to_origin(arm_tag:ArmTag)
    - 你需要对每一个 self.move 进行注释，说明这个动作/组合动作是做什么的，这里的注释要**使用英文**。

举例：
- 没加注释前：
def play_once(self):
    arm_tag = ArmTag('right' if self.object.get_pose().p[0] > 0 else 'left')

    self.move(self.grasp_actor(self.object, arm_tag=arm_tag, pre_grasp_dis=0.1))
    self.move(self.move_by_displacement(arm_tag=arm_tag, z=0.06))
    
    displaystand_pose = self.displaystand.get_functional_point(0)
    
    self.move(self.place_actor(self.object, arm_tag=arm_tag, target_pose=displaystand_pose, constrain='free', pre_dis=0.07))

    self.info['info'] = {'{A}': f"{self.selected_modelname}/base{self.selected_model_id}", '{B}': f"074_displaystand/base{self.displaystand_id}", '{a}': f'{arm_tag}'} 
    return self.info

- 加注释后：
```python
def play_once(self):
    arm_tag = ArmTag('right' if self.object.get_pose().p[0] > 0 else 'left')
    
    # Grasp the object
    self.move(self.grasp_actor(self.object, arm_tag=arm_tag, pre_grasp_dis=0.1))
    # Move up
    self.move(self.move_by_displacement(arm_tag=arm_tag, z=0.06))
    
    # Get display stand's functional point as target pose
    displaystand_pose = self.displaystand.get_functional_point(0)
    
    # Place the object on the display stand
    self.move(self.place_actor(self.object, arm_tag=arm_tag, target_pose=displaystand_pose, constrain='free', pre_dis=0.07))

    self.info['info'] = {'{A}': f"{self.selected_modelname}/base{self.selected_model_id}", '{B}': f"074_displaystand/base{self.displaystand_id}", '{a}': f'{arm_tag}'} 
    return self.info
```

约束：你不能对代码内容进行任何修改，只需要添加注释即可。

现在，开始你的工作！你只需要输出代码块即可。
"""

import os
import time
import traceback
from openai import OpenAI


def parse(source, max_try=5, verbose=True):
    """AI 解析代码"""
    os.environ["all_proxy"] = ""
    os.environ["http_proxy"] = ""
    os.environ["https_proxy"] = ""
    client = OpenAI(api_key="", base_url="")
    start, try_times = time.time(), 0
    while try_times < max_try:
        try_times += 1
        try:
            response = client.chat.completions.create(
                model="deepseek-v3-250324",
                messages=[
                    {
                        "role": "system",
                        "content": system_prompt
                    },
                    {
                        "role": "user",
                        "content": source
                    },
                ],
                stream=True,
            )
            on_thinking, thinking, answering = False, "", ""
            process_count, process = 0, ["|", "/", "-", "\\"]
            for chunk in response:
                content = chunk.choices[0].delta.content
                if content is None:
                    continue
                if hasattr(chunk.choices[0].delta, "reasoning_content"):
                    thinking += chunk.choices[0].delta.reasoning_content
                else:
                    if content == "<think>":
                        on_thinking = True
                    elif content == "</think>":
                        on_thinking = False
                    elif on_thinking:
                        thinking += content
                    else:
                        answering += content

                if verbose:
                    process_count = (process_count + 1) % 4
                    process_show = (thinking + answering)[-50:].replace("\n", "")
                    print(f'\r  {" "*100}', end="")
                    print(f"\r{process[process_count]} {process_show}", end="", flush=True)
        except SyntaxError:
            print(traceback.format_exc())
            break
        except:
            print(traceback.format_exc())
            continue

        result = re.search(r"```python\n([\s\S]*?)\n```", answering, re.S)
        if result is not None:
            if verbose:
                print(
                    f"cost {time.time()-start:.2f}s, try {try_times} time(s)",
                    flush=True,
                )
            return result.group(1)
    return None


def main(file_path, max_try=5, verbose=True):
    try_count = 0
    while try_count < max_try:
        try_count += 1
        # Step 1: 提取类中的方法
        method_source = get_method_source(file_path, "play_once")

        # Step 2: 调用 AI 解析代码
        processed_source = parse(method_source, max_try=5, verbose=verbose)

        # Step 3: 比较两个方法
        if compare_functions(method_source, processed_source):
            replace_method_in_file_with_comments(file_path, "play_once", processed_source)
            break

    if try_count >= max_try:
        with open("error.log", "a", encoding="utf-8") as f:
            f.write(f"Error processing {file_path}: Exceeded maximum retries.\n")


from threading import Thread
from pathlib import Path


def batch(batch_size=5, root="./envs"):
    name_list = [
        "beat_block_hammer",
        "blocks_ranking_rgb",
        "blocks_ranking_size",
        "dump_bin_bigbin",
        "grab_roller",
        "lift_pot",
        "move_can_pot",
        "move_playingcard_away",
        "move_stapler_pad",
        "place_a2b_left",
        "place_a2b_right",
        "place_bread_basket",
        "place_bread_skillet",
        "place_can_basket",
        "place_cylinder_box",
        "place_fan",
        "place_medicine_spot",
        "place_mouse_pad",
        "place_object_scale",
        "place_object_stand",
        "place_phone_stand",
        "place_remote_storage",
        "place_object_basket",
        "put_bottles_dustbin",
        "rotate_qrcode",
        "shake_bottle",
        "shake_bottle_horizontally",
        "place_shoe",
        "slide_mouse_pad",
        "stamp_seal",
        "handover_block",
        "stack_blocks_three",
        "stack_blocks_two",
        "adjust_bottle",
        "stack_bowls_three",
        "stack_bowls_two",
        "click_alarmclock",
        "click_bell",
        "place_container_plate",
        "pick_diverse_bottles",
        "pick_dual_bottles",
        "place_dual_shoes",
        "place_empty_cup",
        "place_object_into_plasticbox",
    ]
    process_list = []
    for file in name_list:
        file = Path(root) / f"{file}.py"
        if file.exists():
            process_list.append(file)
        else:
            print(f"WARNNING: {file.name} not exists!")

    # Create and start a thread for each file
    threads = []
    finish_count, total_count = 0, len(process_list)
    for file in process_list:
        thread = Thread(target=main, args=(file, 5, False))
        thread.start()
        threads.append([file, thread])
        while len(threads) >= batch_size:
            for t in threads:
                if not t[1].is_alive():
                    threads.remove(t)
                    finish_count += 1
                    print(
                        f"[{finish_count:>3d}/{total_count:03d}] files processed. new finish:",
                        t[0].name,
                        flush=True,
                    )
            time.sleep(0.1)

    # Wait for all threads to complete
    while len(threads) > 0:
        for t in threads:
            if not t[1].is_alive():
                threads.remove(t)
                finish_count += 1
                print(
                    f"[{finish_count:>3d}/{total_count:03d}] files processed. new finish:",
                    t[0].name,
                    flush=True,
                )
        time.sleep(0.1)


if __name__ == "__main__":
    batch()
