import os
import re
import argparse
from datetime import datetime
import random

from code_gen.spec_gen_with_skill import spec_gen_with_skill_save
from code_gen.code_gen_with_skill import (
    code_gen_with_skill_save,
    realworld_code_gen_with_skill_save,
)
from code_gen.exploration_gen_with_skill import rlbench_exploration_gen_with_skill_save
from code_gen.plan_gen import solve_pddl
from verification.verification_code import verify_code_with_spec
from validation.total_validation import validate_skill


def read_file(file_path):
    with open(file_path, "r") as file:
        return file.read()


def write_file(file_path, content):
    with open(file_path, "w") as file:
        file.write(content)


def write_log(log_file, message):
    with open(log_file, "a") as log:
        log.write(f"{message}\n")


def split_observation_and_goal(problem_pddl_text):
    idx = problem_pddl_text.find("(:goal")
    if idx != -1:
        observation_part = problem_pddl_text[:idx]
    else:
        observation_part = problem_pddl_text
    return observation_part


def parse_generated_code_for_skills_rlbench(code_text):
    pattern = r"obs\s*,\s*reward\s*,\s*done\s*=\s*(\w+)\(([^)]*)\)"
    matches = re.findall(pattern, code_text)

    skill_call_list = []
    for match in matches:
        skill_name = match[0]
        args_str = match[1]
        skill_call_code = f"{skill_name}({args_str})"
        skill_call_list.append({"skill_name": skill_name, "skill_call_code": skill_call_code})
    return skill_call_list


def parse_generated_code_for_skills_realworld(code_text):
    mode_to_skill = {
        0: "_execute_pick_and_place",
        1: "execute_pick",
        2: "execute_place",
        3: "execute_push",
        4: "execute_pull",
        5: "execute_sweep",
        6: "execute_rotate",
        7: "execute_go",
        8: "go_to_ready_pose",
        9: "execute_gripper",
        10: "execute_gripper",
    }

    pattern = r"(robot\.execute_movement\([^)]*\))"
    skill_call_list = []

    for m in re.finditer(pattern, code_text):
        full_call = m.group(1)
        args = full_call[len("robot.execute_movement("):-1]

        mode_match = re.search(r"mode\s*=\s*(\d+)", args)
        if not mode_match:
            continue
        mode = int(mode_match.group(1))
        skill_name = mode_to_skill.get(mode)
        if not skill_name:
            continue

        skill_call_list.append({
            "skill_name": skill_name,
            "skill_call_code": full_call,
        })

    return skill_call_list


def get_skill_def_code(skill_name, skill_code_text):
    pattern = rf"(def\s+{skill_name}\(.*?)(?=def\s+|$)"
    matches = re.findall(pattern, skill_code_text, flags=re.DOTALL)
    if matches:
        return matches[0].strip()
    else:
        return ""
    
freeze_range_map = {
    # Example:
    # "switch_open_drawer_1": (1, 3),
    # "move_item_1": (2, 5),
}



# ------------------------------------------------------------
# Pipeline entry
# ------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(description="Run the 4‑step pipeline for tasks.")
    parser.add_argument(
        "--env",
        nargs="+",
        required=True,
        help="Which environment(s) to run. Options: rlbench, realworld",
    )
    parser.add_argument(
        "--type",
        type=int,
        choices=[1, 2, 3, 4, 5, 6, 7],
        required=True,
        help="Type of task to run (1, 2, 3, 4, 5, or 6).",
    )
    parser.add_argument(
        "--obs",
        type=int,
        choices=[1, 2, 3, 4],
        default=3,
        help="Observation‑set to use: 1=high+low, 2=high+low+oracle, 3=all",
    )
    args = parser.parse_args()

    rlbench_tasks = ["open_drawer","open_drawer_2"]
    realworld_tasks = ["long_horizon_1"]


    valid_envs = ["rlbench", "realworld"]
    for e in args.env:
        if e not in valid_envs:
            parser.error(f"Unknown environment: {e}.")

    envs_to_run = args.env
    selected_type = args.type
    seed_count = 2

    # --------------------------------------------------------
    # New task‑type behaviour tables
    # --------------------------------------------------------
    type_threshold_map = {1: -0.05, 2: -0.3, 3: -0.27, 4: -0.27, 5: -0.3, 6: -0.4}

    type_to_steps = {
        1: ["spec", "code"],
        2: ["spec", "code", "verify"],
        3: ["spec", "code", "verify", "validate"],
        4: ["spec", "code", "verify", "validate"],  # CSC‑only validation
        5: ["spec", "code", "verify", "validate"],  # LC‑only validation
        6: ["spec", "code", "codesift"],  # Using codesift validation
        7: ["baseline"]
    }

    log_dir = "Input_Your_Path"
    os.makedirs(log_dir, exist_ok=True)
    log_filename = datetime.now().strftime("%Y%m%d_%H%M%S_log.txt")
    log_file = os.path.join(log_dir, log_filename)

    write_log(
        log_file,
        f"Pipeline started with envs: {envs_to_run} and task type: {selected_type}",
    )
    obs_map = {
        1: ["high", "low"],
        2: ["high", "low", "oracle"],
        3: ["high", "middle", "low", "oracle"],
        4: ["oracle"]
    }

    observation_types = obs_map[args.obs]

    # --------------------------------------------------------
    # Iterate over environments
    # --------------------------------------------------------
    for env_name in envs_to_run:
        if env_name == "rlbench":
            base_path = "Input_Your_Path"
            task_list = rlbench_tasks
            available_skills = ["pick", "place", "move", "rotate", "pull"]
        else:
            base_path = "Input_Your_Path"
            task_list = realworld_tasks
            available_skills = [
                "execute_pick",
                "execute_place",
                "execute_push",
                "execute_pull",
                "execute_sweep",
                "execute_rotate",
                "execute_go",
                "execute_gripper",
                "execute_push_swtich -> push"
            ]

        expanded_task_list = [
            (f"{task}_type{selected_type}", selected_type) for task in task_list
        ]

        for task_name, task_type in expanded_task_list:
            # ---------------------------------------------
            # Determine validation‑mode for this task type
            # ---------------------------------------------
            if task_type == 4:
                validation_mode = "csc"  # CSC‑only
            elif task_type == 5:
                validation_mode = "lc"  # LC‑only
            else:
                validation_mode = "both"  # default: both checks

            base_dir = os.path.join(base_path, task_name)

            domain_pddl_path = os.path.join(base_dir, "pddl", "domain.pddl")
            domain_pddl = read_file(domain_pddl_path)
            exploration_knowledge_path = os.path.join(
                base_dir, "pddl", "exploration_knowledge.pddl"
            )
            exploration_knowledge = read_file(exploration_knowledge_path)

            # ------------- load instruction list -------------
            instruction_path = os.path.join(base_dir, "instruction.py")
            instruction_namespace = {}
            with open(instruction_path, "r") as file:
                exec(file.read(), instruction_namespace)
            instructions = instruction_namespace.get("instructions", [])
            if not instructions:
                write_log(log_file, f"[{env_name}] No instructions found in {instruction_path}")
                continue

            # ------------- load skill code -------------
            skill_code_path = os.path.join(base_dir, "code", "skill_code.py")
            skill_code_text = read_file(skill_code_path)

            # =============================================
            # Iterate over observation types / seeds
            # =============================================
            for obs_type in observation_types:
                if obs_type == "oracle":
                    oracle_obs_path = os.path.join(base_dir, "init_obs_oracle.txt")
                    oracle_obs_full_text = read_file(oracle_obs_path)
                    observation_text_for_all = split_observation_and_goal(oracle_obs_full_text)
                else:
                    obs_1_path = os.path.join(base_dir, f"init_obs_{obs_type}_1.txt")
                    obs_1_text = read_file(obs_1_path)
                    observation_text_for_all = split_observation_and_goal(obs_1_text)

                object_list_path = os.path.join(base_dir, f"object_list_position_{obs_type}.txt")
                object_list = read_file(object_list_path) if os.path.exists(object_list_path) else ""

                for seed in range(1, seed_count + 1):
                    # -----------------------------------------
                    # Per‑instruction processing
                    # -----------------------------------------
                    for idx, instruction in enumerate(instructions, start=1):
                        if obs_type == "oracle":
                            chosen_path = os.path.join(base_dir, "init_obs_oracle.txt")
                        else:
                            chosen_path = (
                                os.path.join(base_dir, f"init_obs_{obs_type}_1.txt")
                                if random.random() < 0.66
                                else os.path.join(base_dir, f"init_obs_{obs_type}_2.txt")
                            )
                        problem_pddl_text = read_file(chosen_path)

                        # ---------- plan generation ----------
                        plan_success, plan_result = solve_pddl(
                            domain_pddl, problem_pddl_text
                        )
                        plan_content = plan_result if plan_success else ""

                        # ---------- Spec generation ----------
                        if "spec" in type_to_steps[task_type]:
                            spec_gen_with_skill_save(
                                env_name,
                                task_name,
                                domain_pddl,
                                observation_text_for_all,
                                instruction,
                                idx,
                                available_skills,
                                skill_code_text,
                                obs_type,
                                seed,
                                plan_content=plan_content,
                            )
                            write_log(
                                log_file,
                                f"[{env_name}] Spec generation done. (Instruction #{idx}, obs={obs_type}, seed={seed})",
                            )

                        # ---------- Code generation & verification ----------
                        spec_path = os.path.join(
                            base_dir, "spec", f"exe_spec_{obs_type}_seed{seed}_{idx}.json"
                        )
                        if not os.path.exists(spec_path):
                            continue  # no spec => skip

                        spec_text = read_file(spec_path)
                        feedback = "none"
                        prev_code = "none"
                        max_attempts = 3 if "verify" in type_to_steps[task_type] else 1

                        for attempt in range(max_attempts):
                            if env_name == "rlbench":
                                response = code_gen_with_skill_save(
                                    task_name,
                                    domain_pddl,
                                    observation_text_for_all,
                                    instruction,
                                    spec_text,
                                    skill_code_text,
                                    available_skills,
                                    object_list,
                                    feedback,
                                    prev_code,
                                    exploration_knowledge,
                                    obs_type,
                                    seed,
                                    validated_until_skill_index=0,
                                )
                            else:
                                response = realworld_code_gen_with_skill_save(
                                    task_name,
                                    domain_pddl,
                                    observation_text_for_all,
                                    instruction,
                                    spec_text,
                                    skill_code_text,
                                    available_skills,
                                    object_list,
                                    feedback,
                                    prev_code,
                                    obs_type,
                                    seed,
                                    validated_until_skill_index=0,
                                )

                            code_dir = os.path.join(base_dir, "code")
                            os.makedirs(code_dir, exist_ok=True)
                            code_file_path = os.path.join(
                                code_dir, f"exe_code_{obs_type}_seed{seed}_{idx}.py"
                            )
                            write_file(code_file_path, response)

                            write_log(
                                log_file,
                                f"[{env_name}] Code generation attempt {attempt + 1} (Instruction #{idx}, obs={obs_type}, seed={seed})",
                            )

                            # ------------- Verification -------------
                            if "verify" in type_to_steps[task_type]:
                                generated_code = read_file(code_file_path)
                                passed, feedback = verify_code_with_spec(
                                    task_name,
                                    idx,
                                    domain_pddl,
                                    observation_text_for_all,
                                    instruction,
                                    spec_text,
                                    generated_code,
                                )
                                write_log(
                                    log_file,
                                    f"[{env_name}] Verification attempt {attempt + 1} (instruction #{idx}, obs={obs_type}, seed={seed}): {'Passed' if passed else 'Failed'}",
                                )
                                if passed:
                                    break
                                prev_code = generated_code

                    # ------------------------------------------------
                    # Validation phase (after all instructions)
                    # ------------------------------------------------

                    if "validate" in type_to_steps[task_type]:
                        # print("1")
                        code_dir = os.path.join(base_dir, "code")
                        val_log_path = os.path.join(
                            code_dir, f"validation_log_{obs_type}_seed{seed}.txt"
                        )
                        open(val_log_path, "w").close()

                        threshold = type_threshold_map.get(task_type, -0.5)

                        # Base observation for validation
                        if obs_type == "oracle":
                            base_obs_path = os.path.join(base_dir, "init_obs_oracle.txt")
                        else:
                            base_obs_path = os.path.join(base_dir, f"init_obs_{obs_type}_1.txt")
                        base_obs_text = read_file(base_obs_path)
                        current_problem_pddl_text = base_obs_text
                        current_observation = split_observation_and_goal(base_obs_text)

                        exploration_updates_limit = 2
                        exploration_count = 0
                        for idx, instruction in enumerate(instructions, start=1):
                            exploration_count = 0
                            code_file_path = os.path.join(
                                code_dir, f"exe_code_{obs_type}_seed{seed}_{idx}.py"
                            )
                            if not os.path.exists(code_file_path):
                                write_log(
                                    log_file,
                                    f"[{env_name}] Missing code for validation. (instruction #{idx})",
                                )
                                continue

                            spec_path = os.path.join(
                                base_dir, "spec", f"exe_spec_{obs_type}_seed{seed}_{idx}.json"
                            )
                            if not os.path.exists(spec_path):
                                continue

                            spec_text = read_file(spec_path)
                            max_iterations = 4
                            trial = 0
                            validated_until_skill_index = 0
               
                            while trial < max_iterations:
                                trial += 1
                                generated_code_text = read_file(code_file_path)

                                if env_name == "rlbench":
                                    skill_call_list = parse_generated_code_for_skills_rlbench(generated_code_text)
                                else:
                                    skill_call_list = parse_generated_code_for_skills_realworld(generated_code_text)

                            

                                pass_results = []
                                exploration_triggered = False

                                for sc_idx, sc in enumerate(skill_call_list):
                                  
                                    if sc_idx < validated_until_skill_index:
                                        pass_results.append(1)
                                        continue

                                    skill_name = sc["skill_name"]
                                    skill_call_code = sc["skill_call_code"]
                                    if env_name == "rlbench":
                                        skill_def_code = get_skill_def_code(skill_name, skill_code_text)
                                    else:
                                        skill_def_code = ""


                                    (
                                        csc_pass,
                                        lc_pass,
                                        updated_obs,
                                        csc_fb,
                                        lc_fb,
                                        exploration_needed,
                                        llm_score,
                                        confidence,
                                    ) = validate_skill(
                                        skill_def_code,
                                        skill_call_code,
                                        domain_pddl,
                                        base_dir,
                                        current_observation,
                                        skill_name,
                                        [],
                                        instruction,
                                        object_list,
                                        validation_mode=validation_mode,
                                        threshold=threshold,
                                    )

                                    current_observation = updated_obs
                                    pass_results.append(int(csc_pass and lc_pass))

                                    if not (csc_pass and lc_pass):
                                        write_log(
                                            log_file,
                                            f"[{env_name}] [Inst {idx}] Validation FAILED (skill: {skill_name})",
                                        )
                                    else:
                                        write_log(
                                            log_file,
                                            f"[{env_name}] [Inst {idx}] Validation PASSED (skill: {skill_name})",
                                        )

                                    if exploration_needed:
                                        write_log(
                                            log_file,
                                            f"[{env_name}] Exploration needed (skill: {skill_name})",
                                        )
                                        feedback = f"{csc_fb}\n{lc_fb}"
                                        exploration_code = rlbench_exploration_gen_with_skill_save(
                                            domain_pddl,
                                            current_observation,
                                            skill_def_code,
                                            available_skills,
                                            object_list,
                                            feedback,
                                            exploration_knowledge,
                                        )
                                        exploration_dir = os.path.join(base_dir, "exploration")
                                        os.makedirs(exploration_dir, exist_ok=True)
                                        exploration_file_path = os.path.join(
                                            exploration_dir,
                                            f"exploration_{obs_type}_inst{idx}_{trial}_seed{seed}_type{task_type}.py",
                                        )
                                        write_file(exploration_file_path, exploration_code)

                                        exploration_count += 1
                                        if exploration_count <= exploration_updates_limit:
                                            after_exploration_obs_path = os.path.join(
                                                base_dir,
                                                f"init_obs_after_exploration_{exploration_count}.txt",
                                            )
                                            if os.path.exists(after_exploration_obs_path):
                                                user_updated_observation = read_file(after_exploration_obs_path)
                                                current_problem_pddl_text = user_updated_observation
                                                current_observation = split_observation_and_goal(
                                                    user_updated_observation
                                                )
                                                write_log(
                                                    log_file,
                                                    f"[{env_name}] Observation updated from {after_exploration_obs_path}.",
                                                )
                                                oracle_obj_path = os.path.join(base_dir, "object_list_position_oracle.txt")
                                                object_list = (
                                                    read_file(oracle_obj_path) if os.path.exists(oracle_obj_path) else object_list
                                                )
                                            else:
                                                write_log(
                                                    log_file,
                                                    f"[{env_name}] File not found: {after_exploration_obs_path}. Keeping old observation.",
                                                )
                                        else:
                                            write_log(
                                                log_file,
                                                f"[{env_name}] Exploration update limit exceeded. No more observation_updates.",
                                            )

                                        # Plan & spec regeneration
                                        plan_success, plan_result = solve_pddl(
                                            domain_pddl, current_problem_pddl_text
                                        )
                                        if plan_success:
                                            print(task_name, "plan success")
                                            print(plan_result)
                                            spec_gen_with_skill_save(
                                                env_name,
                                                task_name,
                                                domain_pddl,
                                                current_observation,
                                                instruction,
                                                idx,
                                                available_skills,
                                                skill_code_text,
                                                obs_type,
                                                seed,
                                                plan_content=plan_result,
                                            )
                                        else:
                                            spec_gen_with_skill_save(
                                                env_name,
                                                task_name,
                                                domain_pddl,
                                                current_observation,
                                                instruction,
                                                idx,
                                                available_skills,
                                                skill_code_text,
                                                obs_type,
                                                seed,
                                            )

                                        validated_until_skill_index = sc_idx
                                        prev_code = generated_code_text
                                        feedback = ""

                                        # Code regeneration
                                        if env_name == "rlbench":
                                            new_code = code_gen_with_skill_save(
                                                task_name,
                                                domain_pddl,
                                                current_observation,
                                                instruction,
                                                read_file(
                                                    os.path.join(
                                                        base_dir,
                                                        "spec",
                                                        f"exe_spec_{obs_type}_seed{seed}_{idx}.json",
                                                    )
                                                ),
                                                skill_code_text,
                                                available_skills,
                                                object_list,
                                                feedback,
                                                prev_code,
                                                exploration_knowledge,
                                                obs_type,
                                                seed,
                                                validated_until_skill_index=validated_until_skill_index,
                                            )
                                        else:
                                            plan_content = plan_result if plan_success else ""
                                            new_code = realworld_code_gen_with_skill_save(
                                                task_name,
                                                domain_pddl,
                                                current_observation,
                                                instruction,
                                                read_file(
                                                    os.path.join(
                                                        base_dir,
                                                        "spec",
                                                        f"exe_spec_{obs_type}_seed{seed}_{idx}.json",
                                                    )
                                                ),
                                                skill_code_text,
                                                available_skills,
                                                object_list,
                                                feedback,
                                                prev_code,
                                                obs_type,
                                                seed,
                                                plan_content=plan_content,
                                                validated_until_skill_index=validated_until_skill_index,
                                            )
                                        write_file(code_file_path, new_code)
                                        exploration_triggered = True
                                        break  # break skill loop and retry

                                # end skill loop
                                if exploration_triggered:
                                    continue  # retry whole trial with new code
                                else:
                                    final_confidence = (
                                        sum(pass_results) / len(pass_results) if pass_results else 0.0
                                    )
                                    with open(val_log_path, "a") as vf:
                                        vf.write(
                                            f"=== Instruction {idx}, trial {trial} ===\nConfidence: {final_confidence}\n\n"
                                        )

                                    if final_confidence == 1.0:
                                        break  # validation success for instruction
                                    else:
                                        # Partial failure => regenerate code
                                        validated_until_skill_index = 0
                                        prev_code = generated_code_text
                                        feedback = ""

                                        if env_name == "rlbench":
                                            new_code = code_gen_with_skill_save(
                                                task_name,
                                                domain_pddl,
                                                current_observation,
                                                instruction,
                                                spec_text,
                                                skill_code_text,
                                                available_skills,
                                                object_list,
                                                feedback,
                                                prev_code,
                                                exploration_knowledge,
                                                obs_type,
                                                seed,
                                                validated_until_skill_index=validated_until_skill_index,
                                            )
                                        else:
                                            new_code = realworld_code_gen_with_skill_save(
                                                task_name,
                                                domain_pddl,
                                                current_observation,
                                                instruction,
                                                spec_text,
                                                skill_code_text,
                                                available_skills,
                                                object_list,
                                                feedback,
                                                prev_code,
                                                obs_type,
                                                seed,
                                                validated_until_skill_index=validated_until_skill_index,
                                            )
                                        write_file(code_file_path, new_code)
                                        continue
                        # end instruction loop
                    # end obs loop
                    # ------------------------------------------------
                    # CodeSift validation for type 6
                    # ------------------------------------------------
                    elif "codesift" in type_to_steps[task_type]:
                        from validation.confidence.codesift import codesift_validate
                        
                        code_dir = os.path.join(base_dir, "code")
                        codesift_log_path = os.path.join(
                            code_dir, f"codesift_log_{obs_type}_seed{seed}.txt"
                        )
                        open(codesift_log_path, "w").close()
                        
                        # Base observation
                        if obs_type == "oracle":
                            base_obs_path = os.path.join(base_dir, "init_obs_oracle.txt")
                        else:
                            base_obs_path = os.path.join(base_dir, f"init_obs_{obs_type}_1.txt")
                        base_obs_text = read_file(base_obs_path)
                        current_observation = split_observation_and_goal(base_obs_text)
                        
                        write_log(
                            log_file,
                            f"[{env_name}] Starting CodeSift validation for type 6 (obs={obs_type}, seed={seed})"
                        )
                        
                        for idx, instruction in enumerate(instructions, start=1):
                            code_file_path = os.path.join(
                                code_dir, f"exe_code_{obs_type}_seed{seed}_{idx}.py"
                            )
                            if not os.path.exists(code_file_path):
                                write_log(
                                    log_file,
                                    f"[{env_name}] Missing code for CodeSift validation. (instruction #{idx})",
                                )
                                continue
                            
                            # Get code language based on environment
                            lang = "python"  
                            
                            # Read the code
                            code = read_file(code_file_path)
                            
                            # Combine instruction with domain PDDL and observation for task description
                            task_desc = f"Instruction: {instruction}\n\nDomain PDDL:\n{domain_pddl}\n\nObservation:\n{current_observation}"
                            
                            # Run CodeSift validation
                            is_correct, details = codesift_validate(code, task_desc, lang)
                            
                            # Log results
                            with open(codesift_log_path, "a") as log:
                                log.write(f"=== Instruction {idx} ===\n")
                                log.write(f"Is Correct: {is_correct}\n")
                                log.write(f"Syntax Pass: {details['syntax_pass']}\n")
                                log.write(f"Code Function: {details['code_func']}\n")
                                log.write(f"Similarity: {details['sim']}\n")
                                log.write(f"Differences: {details['diff']}\n\n")
                                
                                if not is_correct and 'refined_code' in details:
                                    log.write("Refined Code Generated:\n")
                                    log.write("-------------------\n")
                                    log.write(f"{details['refined_code']}\n")
                                    log.write("-------------------\n\n")
                                    
                                    # Save the refined code
                                    refined_code_path = os.path.join(
                                        code_dir, f"exe_code_{obs_type}_seed{seed}_{idx}.py"
                                    )
                                    write_file(refined_code_path, details['refined_code'])
                            
                            write_log(
                                log_file,
                                f"[{env_name}] CodeSift validation for instruction #{idx}: {'Passed' if is_correct else 'Failed - Refined code generated'}"
                            )
                    if "baseline" in type_to_steps[task_type]:
                        write_log(log_file, f"[{env_name}] Starting baseline for type 7 (obs={obs_type}, seed={seed})")
                        for idx, instruction in enumerate(instructions, start=1):
                            original_obs = obs_type
                            # Load the original observation
                            obs_file = (os.path.join(base_dir, "init_obs_oracle.txt")
                                        if original_obs == "oracle"
                                        else os.path.join(base_dir, f"init_obs_{original_obs}_1.txt"))
                            observation = split_observation_and_goal(read_file(obs_file))

                            # Initial spec generation
                            spec_gen_with_skill_save(
                                env_name,
                                task_name,
                                domain_pddl,
                                observation,
                                instruction,
                                idx,
                                available_skills,
                                skill_code_text,
                                original_obs,
                                seed,
                            )
                            spec_path = os.path.join(
                                base_dir, "spec", f"exe_spec_{original_obs}_seed{seed}_{idx}.json"
                            )
                            spec_text = read_file(spec_path)

                            # Initial code generation
                            code_response = (
                                code_gen_with_skill_save(
                                    task_name,
                                    domain_pddl,
                                    observation,
                                    instruction,
                                    spec_text,
                                    skill_code_text,
                                    available_skills,
                                    object_list,
                                    feedback="none",
                                    prev_code="none",
                                    exploration_knowledge=None,
                                    obs_type=original_obs,
                                    seed=seed,
                                ) if env_name == "rlbench" else
                                realworld_code_gen_with_skill_save(
                                    task_name,
                                    domain_pddl,
                                    observation,
                                    instruction,
                                    spec_text,
                                    skill_code_text,
                                    available_skills,
                                    object_list,
                                    feedback="none",
                                    prev_code="none",
                                    obs_type=original_obs,
                                    seed=seed,
                                )
                            )
                            code_file = os.path.join(
                                base_dir, "code", f"exe_code_{original_obs}_seed{seed}_{idx}.py"
                            )
                            write_file(code_file, code_response)

                            # Determine freeze_step from skill sequence and hyperparams
                            skill_calls = (parse_generated_code_for_skills_rlbench(code_response)
                                           if env_name == "rlbench"
                                           else parse_generated_code_for_skills_realworld(code_response))
                            num_skills = len(skill_calls)
             
                            freeze_step = random.randint(0, num_skills)
                            write_log(log_file, f"[{env_name}] Instruction {idx}: freezing at skill step {freeze_step}")
                            prev_code_content = code_response

                            # Oracle observation
                            oracle_observation = split_observation_and_goal(
                                read_file(os.path.join(base_dir, "init_obs_oracle.txt"))
                            )

                            # Regenerated spec with frozen_step
                            spec_gen_with_skill_save(
                                env_name,
                                task_name,
                                domain_pddl,
                                oracle_observation,
                                instruction,
                                idx,
                                available_skills,
                                skill_code_text,
                                original_obs,
                                seed
                            )
                            regen_spec_text = read_file(spec_path)

                            # Regenerated code using previous code and frozen_step
                            regen_response = (
                                code_gen_with_skill_save(
                                    task_name,
                                    domain_pddl,
                                    oracle_observation,
                                    instruction,
                                    regen_spec_text,
                                    skill_code_text,
                                    available_skills,
                                    object_list,
                                    feedback="none",
                                    prev_code=prev_code_content,
                                    exploration_knowledge=None,
                                    obs_type=original_obs,
                                    seed=seed,
                                    validated_until_skill_index = freeze_step
                                ) if env_name == "rlbench" else
                                realworld_code_gen_with_skill_save(
                                    task_name,
                                    domain_pddl,
                                    oracle_observation,
                                    instruction,
                                    regen_spec_text,
                                    skill_code_text,
                                    available_skills,
                                    object_list,
                                    feedback="none",
                                    prev_code=prev_code_content,
                                    obs_type=original_obs,
                                    seed=seed,
                                    validated_until_skill_index = freeze_step
                                )
                            )
                            write_file(code_file, regen_response)

                        write_log(log_file, f"[{env_name}] Completed baseline pipeline for type 7")
            
            
            # end obs loop
            write_log(log_file, f"Pipeline completed for task {task_name} at {datetime.now()}")
    write_log(log_file, f"Pipeline fully completed at {datetime.now()}")


if __name__ == "__main__":
    main()

