import os
import json
import re
import numpy as np
from datasets import Dataset


def load_step_data(step_folder):
    all_samples = []
    for filename in os.listdir(step_folder):
        if filename.endswith(".json"):
            filepath = os.path.join(step_folder, filename)
            with open(filepath, "r") as f:
                data = json.load(f)

                initial_prompt = data.get("initial_prompt", "")
                SYSTEM_PROMPT = data.get("system_prompt", "")  # MODIFIED

                for step in data.get("steps", []):
                    if step["step_number"] > 1:
                        user_content = f"Initial Description of Penetration Testing Task: {initial_prompt}"
                        for prev_step in data.get("steps", [])[
                            : step["step_number"] - 1
                        ]:
                            user_content += f"\n===Step {prev_step['step_number']}==="
                            user_content += f"\nThought: {prev_step['thought']}"
                            user_content += f"\nCommand: {prev_step['command']}"
                            user_content += f"\nObservation: {prev_step['result']}"
                        user_content += "\nWhat are the next steps and actions?"
                    else:
                        user_content = f"Initial Description of Penetration Testing Task: {initial_prompt}\nWhat are the next instruction and commands?"

                    content = f"""You will guide the penetration testing process step by step. At each stage, you will receive the context from the previous steps. Based on that information, you must determine the next logical action to advance the assessment. Continue this process until the penetration test is fully completed.

### Response Format (Always follow this format exactly):
=== Step i ===
Thought: [Explain your reasoning and goals here.]
Command: [Write the full command you would run.]

### Previous Steps Context
{user_content}"""

                    assistant_content = f"===Step {step['step_number']}===\nThought: {step['thought']}\nCommand: {step['command']}"

                    sample = {
                        "prompt": [
                            {"role": "system", "content": SYSTEM_PROMPT},
                            {"role": "user", "content": content},
                        ],
                        "answer": assistant_content,
                    }
                    all_samples.append(sample)

    return Dataset.from_list(all_samples)


def main():
    dataset = load_step_data("data/walkthroughs")
    print(dataset)
    print(dataset.features)
    print(dataset.column_names)

    answers = dataset["answer"]
    prompts = dataset["prompt"]
    print(
        answers[:2]
    )  # increasing the index is going to increase the the number of steps visualized


main()
