from stage import Stage

class Pathway:
    def __init__(self, pathway_json, ai_llm):
        self.pathway_json = pathway_json
        self.ai_llm = ai_llm
        self.pathway_complete = False

        self.stages = {}
        for stage_id, stage_data in self.pathway_json.items():
            if stage_id != "pathway_complete":
                self.stages[stage_id] = Stage(stage_data, ai_llm)
        self.current_stage = self.stages["1"]
    
    def next_step(self, user_input):
        next_stage_ui = None
        stage_complete, response, checker_ui = self.current_stage.next_step(user_input)
        if stage_complete:
            print("***Stage Complete***.", self.current_stage.stage_json["name"])
            self.current_stage = self.stages.get(self.current_stage.get_next_stage(), None)
            if self.current_stage is None:
                self.pathway_complete = True
            else:
                next_stage_ui = self.current_stage.get_initial_ui()

        return self.pathway_complete, response, next_stage_ui, checker_ui
    
    def get_current_stage_inital_ui(self):
        return self.current_stage.get_initial_ui()
    
    def get_pathway_complete_ui(self):
        pathway_complete_ui = None
        if "pathway_complete" in self.pathway_json:
            if "ui" in self.pathway_json["pathway_complete"]:
                pathway_complete_ui = self.pathway_json["pathway_complete"]["ui"]
        return pathway_complete_ui