import torch
import re
from collections import defaultdict
import os
from typing import List, Dict, Any, Tuple
from dataclasses import dataclass
from .tensor_helper import TensorHelper, TensorConfig
from verl import DataProto
from verl.utils.tracking import Tracking
import shutil
import requests
import traceback
import sys
from pathlib import Path
import subprocess
import math
from agent.interpreter import Interpreter
import psutil
import wandb
import numpy as np
import copy

@dataclass
class GenerationConfig:
    max_turns: int
    max_start_length: int
    max_prompt_length: int 
    max_response_length: int
    max_obs_length: int
    num_gpus: int
    rollout_n: int
    no_think_rl: bool=False
    working_dir: str = 'working'
    working_temp_dir: str = 'working/tmp'
    working_file_name: str = 'runfile.py'
    csv_folder: str = ''
    format_tb_ipython: bool = False
    trace_back_len: int = 1
    skip_lib: bool = False
    async_rollout_mode: bool = False
    engine_name: str = 'vllm'  # default to vllm, can be changed to other engines like 'huggingface', 'sglang', etc.

class LLMGenerationManager:
    def __init__(
        self,
        tokenizer,
        actor_rollout_wg,
        config: GenerationConfig,
        is_validation: bool = False,
    ):
        self.tokenizer = tokenizer
        self.actor_rollout_wg = actor_rollout_wg
        self.config = config
        self.is_validation = is_validation
        self.interpreter = None
        self.think_str = "<think>"
        self.think_ids = self.tokenizer.encode(self.think_str)
        self.think_ids_tensor = torch.tensor(self.think_ids, dtype=torch.long)

        self.tensor_fn = TensorHelper(TensorConfig(
            pad_token_id=tokenizer.pad_token_id,
            max_prompt_length=config.max_prompt_length,
            max_obs_length=config.max_obs_length,
            max_start_length=config.max_start_length
        ))

    def _batch_tokenize(self, responses: List[str]) -> torch.Tensor:
        """Tokenize a batch of responses."""
        return self.tokenizer(
            responses, 
            add_special_tokens=False, 
            return_tensors='pt', 
            padding="longest"
        )['input_ids']

    def _postprocess_responses(self, responses: torch.Tensor) -> torch.Tensor:
        """Process responses to stop at code operation or answer operation."""
        responses_str = self.tokenizer.batch_decode(
            responses,
            skip_special_tokens=True
        )

        responses_str = [resp.split('</code>')[0] + '</code>' + '<|im_end|>\n'
                 if '</code>' in resp 
                 else resp.split('</answer>')[0] + '</answer>' + '<|im_end|>\n'
                 if '</answer>' in resp 
                 else resp + '<|im_end|>\n'
                 for resp in responses_str]

        if self.config.no_think_rl:
            raise ValueError('stop')
            # if no_think_rl is enabled, only keep action in the str
            actions, _ = self.env.postprocess_predictions(responses_str)
            responses_str=[f"<answer>{envs[idx].ACTION_LOOKUP[action]}</answer>" for idx, action in enumerate(actions)]
            print("RESPONSES:", responses_str)
        responses = self._batch_tokenize(responses_str)
        return responses, responses_str

    def _process_next_obs(self, next_obs: List[str]) -> torch.Tensor:
        """Process next observations from environment."""
        
        next_obs_ids = self.tokenizer(
            next_obs,
            padding='longest',
            return_tensors='pt',
            add_special_tokens=False,  # Prevents adding special tokens
        )['input_ids']

        if next_obs_ids.shape[1] > self.config.max_obs_length:
            print(f"[WARNING] OBSERVATION TOO LONG, CONSIDER CHANGING YOUR CONFIG, {next_obs_ids.shape[1]} & {self.config.max_obs_length}")            
            next_obs_ids = next_obs_ids[:, :self.config.max_obs_length]

        return next_obs_ids

    def _update_rolling_state(self, rollings: DataProto, cur_responses: torch.Tensor, 
                            next_obs_ids: torch.Tensor,
                            need_thinks) -> Dict:
        """Update rolling state with new responses and observations."""
        
        need_thinks = None if self.config.async_rollout_mode else need_thinks

        # Concatenate and handle padding     
        # breakpoint()
        think_ids_expanded = self.think_ids_tensor.unsqueeze(0).expand(rollings.batch['input_ids'].shape[0], -1)

        new_input_ids = self.tensor_fn.concatenate_with_padding([
            rollings.batch['input_ids'],
            cur_responses,
            next_obs_ids,
            think_ids_expanded
        ])
        
        # Create attention mask and position ids
        new_attention_mask = self.tensor_fn.create_attention_mask(new_input_ids)
        new_position_ids = self.tensor_fn.create_position_ids(new_attention_mask)

        # Cut to appropriate length
        effective_len = new_attention_mask.sum(dim=1).max()
        max_len = min(self.config.max_prompt_length, effective_len)

        tensor_batch = {
            'input_ids': new_input_ids[:, -max_len:],
            'position_ids': new_position_ids[:, -max_len:],
            'attention_mask': new_attention_mask[:, -max_len:]
        }
        
        raw_prompts = rollings.non_tensor_batch['raw_prompt']
        new_prompts = []

        cutted_responses = cur_responses[:, -self.config.max_response_length:]

        assistant_responses = self.tokenizer.batch_decode(
            cutted_responses,
            skip_special_tokens=True
        )

        execution_results = self.tokenizer.batch_decode(
            next_obs_ids,
            skip_special_tokens=True
        )
        
        for prompt, assistant_response, execution_result in zip(
            raw_prompts, assistant_responses, execution_results
        ): 
            # if assistant_response:
            #     match = re.search(r"<interpreter>(.*?)</interpreter>", execution_result, re.DOTALL)
            #     if match:
            #         extracted_result = match.group(1).strip()
            #     else:
            #         extracted_result = execution_result
            #     prompt_list = list(prompt)  
            #     new_prompt = prompt_list + [
            #         {"content": assistant_response, "role": "assistant"},
            #         {"content": extracted_result, "role": "user"}
            #     ]
            #     new_prompts.append(new_prompt)
            # else:
            #     prompt_list = list(prompt)
            #     new_prompts.append(prompt_list)
        
            match = re.search(r"<interpreter>(.*?)</interpreter>", execution_result, re.DOTALL)
            if match:
                extracted_result = match.group(1).strip()
            else:
                extracted_result = execution_result
            prompt_list = list(prompt)
            new_prompt = prompt_list + [
                {"content": assistant_response, "role": "assistant"},
                {"content": extracted_result, "role": "user"}
            ]
            new_prompts.append(new_prompt)
            
        non_tensor_batch = {
            "raw_prompt": new_prompts,
        }

        new_rollings = DataProto.from_dict(tensor_batch, non_tensor_batch)
        new_rollings.meta_info.update(rollings.meta_info)
        
        return new_rollings

    def _exec_result_masked_concatenate_with_padding(self, 
            prompt: torch.Tensor, 
            prompt_with_mask: torch.Tensor, 
            response: torch.Tensor, 
            exec_result: torch.Tensor = None,
            need_thinks = None,
            pad_to_left: bool = True
        ) -> torch.Tensor:
        """Concatenate tensors and handle padding. Additionally, create a mask (exec_mask) to cover the execution result block(<result>) if it exists."""
        pad_id = self.tokenizer.pad_token_id
        tensors = [prompt, response]
        tensors_with_mask = [prompt_with_mask, response]
        if exec_result is not None:
            tensors.append(exec_result)
            exec_mask = torch.full(exec_result.size(), pad_id, dtype=exec_result.dtype, device=exec_result.device) # exe mask
            tensors_with_mask.append(exec_mask)

        if need_thinks is not None:
            result = []
            for flag in need_thinks:
                if flag == 1:
                    result.append(self.think_ids_tensor.clone())
                else:
                    result.append(torch.full_like(self.think_ids_tensor, pad_id))

            output = torch.stack(result)
            # print(output)
            tensors.append(output)
            tensors_with_mask.append(output)


        concatenated = torch.cat(tensors, dim=1)  # origin prompt and response
        concatenated_with_mask = torch.cat(tensors_with_mask, dim=1)  # prompt, response, result with mask
        mask = concatenated != pad_id if pad_to_left else concatenated == pad_id
        sorted_indices = mask.to(torch.int64).argsort(dim=1, stable=True)
        padded_tensor = concatenated.gather(1, sorted_indices)
        padded_tensor_with_mask = concatenated_with_mask.gather(1, sorted_indices)

        return padded_tensor, padded_tensor_with_mask

    def _update_right_side(self, right_side: Dict, 
                          cur_responses: torch.Tensor,
                          next_obs_ids: torch.Tensor = None,
                          need_thinks = None
                           ) -> Dict:
        """Update right side state."""

        need_thinks = None if self.config.async_rollout_mode else need_thinks

        if next_obs_ids != None:
            responses, responses_with_info_mask = self._exec_result_masked_concatenate_with_padding(
                    right_side['responses'],
                    right_side['responses_with_info_mask'],
                    cur_responses,
                    next_obs_ids,
                    need_thinks,
                    pad_to_left=False
                )
        else:
            responses, responses_with_info_mask = self._exec_result_masked_concatenate_with_padding(
                    right_side['responses'],
                    right_side['responses_with_info_mask'],
                    cur_responses,
                    pad_to_left=False
                )
        effective_len = self.tensor_fn.create_attention_mask(responses).sum(dim=1).max()
        max_len = min(self.config.max_prompt_length, effective_len)
        
        return {'responses': responses[:, :max_len], 'responses_with_info_mask': responses_with_info_mask[:, :max_len]}
    
    def rollout(self, batch: DataProto) -> DataProto:
        print(f"[DEBUG] gen_batch meta info: {batch.meta_info}")
        if not self.config.async_rollout_mode:
            rollout_batch = self.actor_rollout_wg.generate_sequences(batch)
        else:
            self.actor_rollout_wg.wake_up()
            rollout_batch = self.actor_rollout_wg.generate_sequences(batch)
            self.actor_rollout_wg.sleep()

        return rollout_batch

    def _generate_with_gpu_padding(self, active_batch: DataProto) -> DataProto:
        """
            Wrapper for generation that handles multi-GPU padding requirements.
            if num_gpus <= 1, return self.actor_rollout_wg.generate_sequences(active_batch)
            if active_batch size is not divisible by num_gpus, pad with first sequence
            then remove padding from output
        """
        num_gpus = self.config.num_gpus
        if num_gpus <= 1:
            return self.rollout(active_batch)
            
        batch_size = active_batch.batch['input_ids'].shape[0]
        remainder = batch_size % num_gpus
        
        for key in active_batch.batch.keys():
            active_batch.batch[key] = active_batch.batch[key].long()
        if remainder == 0:
            return self.rollout(active_batch)
        
        # Add padding sequences
        padding_size = num_gpus - remainder
        padded_batch = {}
        padded_non_tensor_batch = {}
        
        for k, v in active_batch.batch.items():
            # Use first sequence as padding template
            pad_sequence = v[0:1].repeat(padding_size, *[1] * (len(v.shape) - 1))
            padded_batch[k] = torch.cat([v, pad_sequence], dim=0)
        
        for key, value in active_batch.non_tensor_batch.items():
            if isinstance(value, (list, np.ndarray)) and len(value) == batch_size:
                if isinstance(value, list):
                    pad_item = value[0]
                    padded_value = value + [pad_item] * padding_size
                elif isinstance(value, np.ndarray):
                    pad_item = value[0:1]  # shape (1, feature_dim)
                    padded_value = np.concatenate([value, np.repeat(pad_item, padding_size, axis=0)], axis=0)
                padded_non_tensor_batch[key] = padded_value
            else:
                # Just copy over fields not tied to batch size
                padded_non_tensor_batch[key] = copy.deepcopy(value)

        padded_active_batch = DataProto.from_dict(padded_batch, padded_non_tensor_batch)
        for key in padded_active_batch.batch.keys():
            padded_active_batch.batch[key] = padded_active_batch.batch[key].long()
        
        padded_active_batch.meta_info.update(active_batch.meta_info)

        # Generate with padded batch
        padded_output = self.rollout(padded_active_batch)

        # Remove padding from output
        trimmed_batch = {k: v[:-padding_size] for k, v in padded_output.batch.items()}
        trimmed_non_tensor_batch = {k: v[:-padding_size] for k, v in padded_output.non_tensor_batch.items()}
        
        # Handle meta_info if present
        if hasattr(padded_output, 'meta_info') and padded_output.meta_info:
            trimmed_meta = {}
            for k, v in padded_output.meta_info.items():
                if isinstance(v, torch.Tensor):
                    trimmed_meta[k] = v[:-padding_size]
                else:
                    trimmed_meta[k] = v
            padded_output.meta_info = trimmed_meta
            
        padded_output.batch = trimmed_batch
        padded_output.non_tensor_batch = trimmed_non_tensor_batch
        return padded_output
    
    def prepare_workspace(self):
        working_path = Path(self.config.working_dir).resolve()

        if not os.path.exists(working_path):
            os.makedirs(working_path, exist_ok=True)

        origin_workdir = os.getcwd()

        os.chdir(str(working_path))

        # copy the csv from local to workspace
        src_folder = self.config.csv_folder 
        dest_folder = working_path / 'data/files'

        self.copy_files(src_folder, dest_folder)

        os.chdir(str(origin_workdir))

    def prepare_interpreter(self, batch_size, timeout = 5):
        if batch_size < self.worker_num:
            self.interpreters.append(Interpreter(timeout=timeout, batch_size=batch_size))
        else:
            interpreter_num = math.floor(batch_size / self.worker_num)
            n = batch_size % self.worker_num
            for i in range(interpreter_num):
                self.interpreters.append(Interpreter(timeout = timeout, batch_size=self.worker_num))
            if n > 0:
                self.interpreters.append(Interpreter(batch_size=n, timeout=timeout))

    def clear_cache(self, activate_list):
        cache_list = [activate_list[i:i + self.worker_num] for i in range(0, len(activate_list), self.worker_num)]
        
        for i, act_list in enumerate(cache_list):
            self.interpreters[i].clear_cache(act_list)
    
    def close_interpreter(self):
        for interpreter in self.interpreters:
            interpreter.close()
        
        self.interpreters = []

    def run_llm_loop(self, gen_batch, initial_input_ids: torch.Tensor) -> Tuple[Dict, Dict]:
        """Run main LLM generation loop."""
        original_left_side = {'input_ids': initial_input_ids[:, -self.config.max_start_length:]}
        original_right_side = {'responses': initial_input_ids[:, []], 'responses_with_info_mask': initial_input_ids[:, []]}

        # concatenate the first loop input with <think>\n
        if not self.config.async_rollout_mode and self.config.engine_name == 'vllm':
            think_ids_expanded = self.think_ids_tensor.unsqueeze(0).expand(gen_batch.batch['input_ids'].shape[0], -1)
            original_right_side['responses'] = torch.cat([original_right_side['responses'], think_ids_expanded], dim=1)
            original_right_side['responses_with_info_mask'] = torch.cat([original_right_side['responses_with_info_mask'], think_ids_expanded], dim=1)

        length = gen_batch.batch['input_ids'].shape[0] * self.config.rollout_n
        active_mask = torch.ones(gen_batch.batch['input_ids'].shape[0], dtype=torch.bool)
        turns_stats = torch.ones(gen_batch.batch['input_ids'].shape[0], dtype=torch.int)
        valid_action_stats = torch.zeros(gen_batch.batch['input_ids'].shape[0], dtype=torch.int)
        valid_exec_stats = torch.zeros(gen_batch.batch['input_ids'].shape[0], dtype=torch.int)
        active_num_list = [active_mask.sum().item()]

        rollings = gen_batch
        if not self.config.async_rollout_mode and self.config.engine_name == 'vllm':
            rollings.batch['input_ids'] = torch.cat([rollings.batch['input_ids'], think_ids_expanded], dim=1)
        rollings.batch['input_ids'] = rollings.batch['input_ids'][:, -self.config.max_prompt_length:]
        rollings.batch['attention_mask'] = self.tensor_fn.create_attention_mask(rollings.batch['input_ids'])
        rollings.batch['position_ids'] = self.tensor_fn.create_position_ids(rollings.batch['attention_mask'])

        self.interpreter = Interpreter(batch_size=length)

        self.prepare_workspace()

        # Main generation loop
        for step in range(self.config.max_turns):
            if not active_mask.sum():
                break
            # cut the paddings in rollings

            print(f"[DEBUG] {step} runs")

            rollings.batch = self.tensor_fn.cut_to_effective_len(
                rollings.batch,
                keys=['input_ids', 'attention_mask', 'position_ids']
            )

            activate_batch_batch = {
                k: v[active_mask] for k, v in rollings.batch.items()
            }
            activate_batch_non_tensor_batch = {
                k: v[active_mask] for k, v in rollings.non_tensor_batch.items()
            }

            rollings_active = DataProto.from_dict(activate_batch_batch, activate_batch_non_tensor_batch)
            # rollings_active = DataProto.from_dict({
            #     k: v[active_mask] for k, v in rollings.batch.items()
            # })

            rollings_active.meta_info.update(rollings.meta_info)
            
            print(f"[DEBUG] generate_sequences start")
            # breakpoint()
            gen_output = self._generate_with_gpu_padding(rollings_active)
            print(f"[DEBUG] generate_sequences end")

            meta_info = gen_output.meta_info

            # decode the answer and extract <code> <answer> or do no action if they don't appear in the response
            # generate new responses_ids by decode responses with skip-special-token and then encode it
            responses_ids, responses_str = self._postprocess_responses(gen_output.batch['responses'])
            # padding the item that active_mask == 0
            responses_ids, responses_str = self.tensor_fn._example_level_pad(responses_ids, responses_str, active_mask)

            print(f"[DEBUG] execute_code start")
            # Execute in environment and process observations
            next_obs, dones, valid_action, is_exec, need_thinks = self.execute_code(
                responses_str, self.tokenizer.pad_token, active_mask
            )
            print(f"[DEBUG] execute_code end")
            
            curr_active_mask = torch.tensor([not done for done in dones], dtype=torch.bool)
            active_mask = active_mask * curr_active_mask
            active_num_list.append(active_mask.sum().item())
            turns_stats[curr_active_mask] += 1
            valid_action_stats += torch.tensor(valid_action, dtype=torch.int)
            valid_exec_stats += torch.tensor(is_exec, dtype=torch.int)

            activate_list = curr_active_mask.tolist()
            self.interpreter.clear_cache(activate_list)
            # self.clear_cache(activate_list)

            # tokenize obs and padding it
            next_obs_ids = self._process_next_obs(next_obs)
            
            # Update states
            rollings = self._update_rolling_state(
                rollings,
                responses_ids,
                next_obs_ids,
                need_thinks
            )

            original_right_side = self._update_right_side(
                original_right_side,
                responses_ids,
                next_obs_ids,
                need_thinks
            )
            
        # final LLM rollout(after think, code, result, finally come out answer)
        if active_mask.sum():

            print(f"[DEBUG] final LLM rollout")
            rollings.batch = self.tensor_fn.cut_to_effective_len(
                rollings.batch,
                keys=['input_ids', 'attention_mask', 'position_ids']
            )

            activate_batch_batch = {
                k: v[active_mask] for k, v in rollings.batch.items()
            }
            activate_batch_non_tensor_batch = {
                k: v[active_mask] for k, v in rollings.non_tensor_batch.items()
            }

            rollings_active = DataProto.from_dict(activate_batch_batch, activate_batch_non_tensor_batch)
            # gen_output = self.actor_rollout_wg.generate_sequences(rollings)
            # rollings_active = DataProto.from_dict({
            #     k: v[active_mask] for k, v in rollings.batch.items()
            # })
            
            print(f"[DEBUG] generate_sequences start")
            rollings_active.meta_info.update(rollings.meta_info)
            gen_output = self._generate_with_gpu_padding(rollings_active)
            print(f"[DEBUG] generate_sequences end")

            meta_info = gen_output.meta_info
            responses_ids, responses_str = self._postprocess_responses(gen_output.batch['responses'])
            responses_ids, responses_str = self.tensor_fn._example_level_pad(responses_ids, responses_str, active_mask)

            print(f"[DEBUG] execute_code start")
            # # Execute in environment and process observations
            next_obs, dones, valid_action, is_exec, need_thinks = self.execute_code(
                responses_str, self.tokenizer.pad_token, active_mask
            )
            print(f"[DEBUG] execute_code end")

            # breakpoint()

            curr_active_mask = torch.tensor([not done for done in dones], dtype=torch.bool)
            active_mask = active_mask * curr_active_mask
            active_num_list.append(active_mask.sum().item())
            valid_action_stats += torch.tensor(valid_action, dtype=torch.int)
            valid_exec_stats += torch.tensor(is_exec, dtype=torch.int)

            # tokenize obs and padding it
            next_obs_ids = self._process_next_obs(next_obs)

            original_right_side = self._update_right_side(
                original_right_side,
                responses_ids,
            )
        
        meta_info['turns_stats'] = turns_stats.tolist()
        meta_info['active_mask'] = active_mask.tolist()
        meta_info['valid_action_stats'] = valid_action_stats.tolist()
        meta_info['valid_exec_stats'] = valid_exec_stats.tolist()
        
        print("ACTIVE_TRAJ_NUM:", active_num_list)
        print("ACTIVE_TRAJ_NUM_EXPLANATION: Initialization, First Turn run, Second Turn run, ..., Final run")
        
        return self._compose_final_output(original_left_side, original_right_side, meta_info)

    def _compose_final_output(self, left_side: Dict,
                            right_side: Dict,
                            meta_info: Dict) -> Tuple[Dict, Dict]:
        """Compose final generation output."""
        final_output = right_side.copy()
        final_output['prompts'] = left_side['input_ids']
        
        # Combine input IDs
        final_output['input_ids'] = torch.cat([
            left_side['input_ids'],
            right_side['responses']
        ], dim=1)
        
        # Create attention mask and position ids
        final_output['attention_mask'] = torch.cat([
            self.tensor_fn.create_attention_mask(left_side['input_ids']),
            self.tensor_fn.create_attention_mask(final_output['responses'])
        ], dim=1)
        final_output['info_mask'] = torch.cat([
            self.tensor_fn.create_attention_mask(left_side['input_ids']),
            self.tensor_fn.create_attention_mask(final_output['responses_with_info_mask'])
        ], dim=1)
        
        final_output['position_ids'] = self.tensor_fn.create_position_ids(
            final_output['attention_mask']
        )
        
        final_output = DataProto.from_dict(final_output)
        final_output.meta_info.update(meta_info)
        
        return final_output

    def execute_code(self, predictions: List[str], pad_token: str, active_mask=None) ->List[str]:
        """
        Execute predictions across multiple environments.
        NOTE: the function is the actual `step` function in the environment
        NOTE penalty_for_invalid is not included in observation shown to the LLM
        
        Args:
            predictions: List of action responses
            pad_token: Token to use for padding
            
        Returns:
            List of observation strings
        """
        cur_actions, contents = self.postprocess_predictions(predictions)
        next_obs, dones, valid_action, is_exec, need_thinks = [], [], [], [], []
        
        generation_code = [content for action, content in zip(cur_actions, contents) if action == 'code']
        exec_results = self._execute(generation_code)
        assert len(exec_results) == sum([1 for action in cur_actions if action == 'code'])

        for i, (action, active) in enumerate(zip(cur_actions, active_mask)):
            
            if not active:
                next_obs.append('')
                dones.append(1)
                valid_action.append(0)
                is_exec.append(0)
                need_thinks.append(0)
            else:
                if action == 'answer':
                    next_obs.append('')
                    dones.append(1)
                    valid_action.append(1)
                    is_exec.append(0)
                    need_thinks.append(0)
                elif action == 'code':
                    next_obs.append(f'<|im_start|>user\n<interpreter>\n{exec_results.pop(0).strip()}\n</interpreter><|im_end|>\n<|im_start|>assistant\n')
                    dones.append(0)
                    valid_action.append(1)
                    is_exec.append(1)
                    need_thinks.append(1)
                else:
                    next_obs.append(f'<|im_start|>user\nYour previous action is invalid. \
If you want to execute the code for the execution result, you should put the code between <code>\n```python and ```\n</code>. \
If you want to give the final answer, you should put the answer between <answer> and </answer>. Please try again.<|im_end|>\n<|im_start|>assistant\n')
                    dones.append(0)
                    valid_action.append(0)
                    is_exec.append(0)
                    need_thinks.append(1)
            
        assert len(exec_results) == 0
            
        return next_obs, dones, valid_action, is_exec, need_thinks

    def postprocess_predictions(self, predictions: List[Any]) -> Tuple[List[int], List[bool]]:
        """
        Process (text-based) predictions from llm into actions and validity flags.
        
        Args:
            predictions: List of raw predictions
            
        Returns:
            Tuple of (actions list, validity flags list)
        """
        actions = []
        contents = []
                
        for i, prediction in enumerate(predictions):
            if isinstance(prediction, str): # for llm output
                pattern = r'<(code|answer)>(.*?)</\1>'
                match = re.search(pattern, prediction, re.DOTALL)
                if match:
                    content = match.group(2).strip()  # Return only the content inside the tags
                    action = match.group(1)
                else:
                    content = ''
                    action = None
            else:
                raise ValueError(f"Invalid prediction type: {type(prediction)}")
            
            actions.append(action)
            contents.append((i, content))
            
        return actions, contents
    
    def extract_files_from_code(self, code):
        # 使用正则表达式查找类似文件路径的字符串
        pattern = r'["\']([^"\']*\.(?:txt|csv|json|yaml|pickle|h5|xml|dat))["\']'
        return re.findall(pattern, code)
    
    def run_code(self, codes):
        code_snippers = [[] for _ in range(len(self.interpreters))]
        results = []
        for pos, code in codes:
            interpreter_n = pos // self.worker_num
            pos_new = pos - (interpreter_n * self.worker_num) 
            code_snippers[interpreter_n].append((pos_new, code))
        
        for i, code_snipper in enumerate(code_snippers):
            result = self.interpreters[i].run_code(code_snipper)
            results.extend(result)

        return results
    
    def _execute(self, codes) ->str:
        """
        Batchified execute for codes.
        Args:
            codes: codes to be executed
        Returns:
            execution results which is concatenated into a string
        """
        results = []

        working_path = Path(self.config.working_dir).resolve()
        working_tmp_path = Path(self.config.working_temp_dir).resolve()
        working_file = str(working_path / self.config.working_file_name)

        # if working_dir doesn't exist, make it
        if not os.path.exists(working_path):
            os.makedirs(working_path, exist_ok=True)

        # if working_dir doesn't exist, make it
        if not os.path.exists(working_tmp_path):
            os.makedirs(working_tmp_path, exist_ok=True)

        origin_workdir = os.getcwd()

        os.chdir(str(working_path))

        raw_results = self.interpreter.run_code(codes)
        # raw_results = self.run_code(codes)

        for result, report in raw_results:
            if isinstance(result, str):
                if result.strip() != "" and report.strip() == "Done":
                    results.append(f"The code run successfully:\n{result}")
                elif report.strip() != "Done" and result.strip() == "":
                    results.append(f"The code run failed:\n{report}")
                elif report.strip() != "" and result.strip() != "":
                    results.append(f"The code run failed:\n{report}\n\nBut we capture part of your code output:\n{result}")
                else:
                    results.append("We couldn't capture the output from your code. Please rewrite your last step code and modify it to explicitly use print() statements to display the values of any variables you want to inspect. Make sure to return the complete and corrected version of the code, and ensure that it can run successfully.")
            else:
                results.append("We couldn't capture the output from your code. Please rewrite your last step code and modify it to explicitly use print() statements to display the values of any variables you want to inspect. Make sure to return the complete and corrected version of the code, and ensure that it can run successfully.")     

        os.chdir(str(origin_workdir))

        return results
    
    def copy_files(self, src_folder, dest_folder):
        if not os.path.exists(dest_folder):
            os.makedirs(dest_folder)
        
        assert os.path.exists(src_folder), f"The src_folder {src_folder} not exist!"

        for filename in os.listdir(src_folder):
            src_file = os.path.join(src_folder, filename)
            
            if os.path.isfile(src_file):
                dest_file = os.path.join(dest_folder, filename)

                if os.path.exists(dest_file):
                    continue
                shutil.copy(src_file, dest_file)
    
    def exception_summary(self, e, working_dir, exec_file_name, format_tb_ipython):
        """Generates a string that summarizes an exception and its stack trace (either in standard python repl or in IPython format)."""
        if format_tb_ipython:
            import IPython.core.ultratb

            # tb_offset = 1 to skip parts of the stack trace in weflow code
            tb = IPython.core.ultratb.VerboseTB(tb_offset=1, color_scheme="NoColor")
            tb_str = str(tb.text(*sys.exc_info()))
        else:
            tb_lines = traceback.format_exception(e)
            # skip parts of stack trace in weflow code
            tb_str = "".join(
                [l for l in tb_lines if "importlib" not in l]
            )

        # replace whole path to file with just filename (to remove agent workspace dir)
        tb_str = tb_str.replace(str(working_dir + '/' + exec_file_name), exec_file_name)

        exc_info = {}
        if hasattr(e, "args"):
            exc_info["args"] = [str(i) for i in e.args]
        for att in ["name", "msg", "obj"]:
            if hasattr(e, att):
                exc_info[att] = str(getattr(e, att))

        tb = traceback.extract_tb(e.__traceback__)
        exc_stack = [(t.filename, t.lineno, t.name, t.line) for t in tb]

        return tb_str, e.__class__.__name__, exc_info, exc_stack
    
    def extract_error_details(self, error_string):
        pattern = re.compile(r'File "(?P<file>.*?)", line (?P<line>\d+), in (?P<module>.*?)\n\s*(?P<code>.*?)\n')
        
        matches = pattern.finditer(error_string)
        
        error_details = []
        for match in matches:
            details = match.groupdict()  
            details['line'] = int(details['line']) 
            error_details.append(details)
        
        return error_details
    
    
    def exception_summary_str(self, tb_str, working_dir, exec_file_name):
        """Generates a string that summarizes an exception and its stack trace"""

        if "timeout_decorator.timeout_decorator.TimeoutError: 'Timed Out'" in tb_str:
            err_msg = "The execution time of your code has exceeded timeout. Check if there are any dead loops in your code that may cause timeout"
        elif 'Please do not use forbidden libs in your code' in tb_str:
            err_msg = "Please don't use the forbidden libs in your code such as 'nltk' or 'nltk.corpus' which will cause an unexpected error"
        else:
            err_msg = "Traceback (most recent call last):\n"

        tb_str = tb_str.replace(str(working_dir + '/' + exec_file_name), exec_file_name)


        error_pattern = r'(\w+Error): (.+?)(?=\n\s+\w|$)'  


        match_class = re.search(error_pattern + r'(?!.*' + error_pattern + ')', tb_str, re.DOTALL)

        if not match_class:
            exc_info = ''
        else:
            error_class = match_class.group(1)  # ZeroDivisionError
            error_message = match_class.group(2)  # division by zero
            error_message = error_message.split("The above exception was the direct cause of the following exception:")[0].strip()  # 去掉多余的部分

            exc_info = f"{error_class}: {error_message}"

        exc_stack = []
        error_details = self.extract_error_details(tb_str)

        # 填充 exc_stack
        for err in error_details:
            if self.config.skip_lib:
                if 'site-packages' in err['file']:
                    continue
                else:
                    exc_stack.append((err['file'], err['line'], err['module'], err['code']))
            else:
                exc_stack.append((err['file'], err['line'], err['module'], err['code']))


        # print(f"[DEBUG] len(exc_stack): {len(exc_stack)}")

        if len(exc_stack) > 0:
            result_len = self.config.trace_back_len * -1
            concat_result = "\n".join([f"File {filename}, Line: {lineno}, in {module}\n{code}"
                                    for filename, lineno, module, code in exc_stack[result_len:]])

            err_msg = err_msg + concat_result + "\n"
        else:
            err_msg = err_msg + tb_str + "\n"

        if exc_info:
            err_msg += exc_info

        # print(f"[DEBUG] err_msg: {err_msg}")

        return err_msg