from env.environment import CodeGuardianEnv
from env.models import Action

env = CodeGuardianEnv()
obs = env.reset()

done = False

step_cycle = ["detect", "risk", "fix"]
step_index = 0

while not done:
    print("\n--- NEW STEP ---")
    print("Code:", obs.code)

    current_action = step_cycle[step_index]

    if current_action == "detect":
        action = Action(
            action_type="detect",
            detected_issues=["wrong operator"],
            risk_level=None,
            suggested_fix=None
        )

    elif current_action == "risk":
        action = Action(
            action_type="risk",
            detected_issues=None,
            risk_level="low",
            suggested_fix=None
        )

    elif current_action == "fix":
        action = Action(
            action_type="fix",
            detected_issues=None,
            risk_level=None,
            suggested_fix="return a + b"
        )

    obs, reward, done, info = env.step(action)

    print("Action:", current_action)
    print("Reward:", reward)
    print("Info:", info)

    step_index = (step_index + 1) % 3