from client import MyEnv
from models import MyAction


def main():

    with MyEnv(base_url="http://localhost:8000").sync() as client:

        # 🔁 RESET
        result = client.reset()
        obs = result.observation

        print("\n================ RESET ================")
        print("Symptom:", obs.echoed_message)
        print("Available actions:")
        print_actions(obs.available_actions)
        print("=======================================\n")

        done = False
        step_counter = 0

        while not done:

            print(f"\n--- Step {step_counter} ---")

            # 🧠 PHASE 1: DIAGNOSE
            if "diagnose" in obs.available_actions and not obs.history:
                action_obj = obs.available_actions["diagnose"][0]

            elif "diagnose" in obs.available_actions and len(obs.history) < 2:
                action_obj = obs.available_actions["diagnose"][len(obs.history)]

            # 🔧 PHASE 2: RECTIFY
            elif "rectify" in obs.available_actions:
                action_obj = obs.available_actions["rectify"][0]

            else:
                print("⚠️ No valid actions left")
                break

            action = MyAction(message=action_obj["id"])

            print(f"> Action: {action_obj['id']}")
            print(f"  ↳ {action_obj['description']} (cost={action_obj['cost']})")

            # 🚀 STEP
            result = client.step(action)
            obs = result.observation

            # 📊 PRINT STATE
            print(f"Observation: {obs.echoed_message}")
            print(f"Reward: {result.reward}")
            print(f"Total Cost: {obs.total_cost}")
            print(f"History length: {len(obs.history)}")

            # 🧠 show last step
            if obs.history:
                last = obs.history[-1]
                print(f"Last Step Type: {last['type']}")
                print(f"Last Action: {last['action']}")

            print("Available actions:")
            print_actions(obs.available_actions)
            print("-----------------------------------")

            done = result.done
            step_counter += 1

        print("\n✅ Episode finished")
        print("Final cost:", obs.total_cost)
        print("Total steps:", step_counter)
        print("===================================\n")


def print_actions(actions):
    for category, items in actions.items():
        print(f"\n[{category.upper()}]")
        for a in items:
            print(f"- {a['id']}: {a['description']} (cost={a['cost']}, reward={a['reward']})")


if __name__ == "__main__":
    main()