"""Pipeline stage: screen user prompt for quality before plan generation."""
from llama_index.core.llms.llm import LLM
from worker_plan_internal.plan.run_plan_pipeline import PlanTask
from worker_plan_internal.diagnostics.screen_planning_prompt import ScreenPlanningPrompt
from worker_plan_api.filenames import FilenameEnum
from worker_plan_api.plan_file import PlanFile
from worker_plan_internal.plan.nodes.initial_plan_raw import InitialPlanRawTask


class ScreenPlanningPromptTask(PlanTask):
    """Flag prompts as UNUSABLE when there is high confidence the prompt is garbage."""
    def requires(self):
        return self.clone(InitialPlanRawTask)

    def output(self):
        return {
            'raw': self.local_target(FilenameEnum.SCREEN_PLANNING_PROMPT_RAW),
            'markdown': self.local_target(FilenameEnum.SCREEN_PLANNING_PROMPT_MARKDOWN)
        }

    def run_with_llm(self, llm: LLM) -> None:
        plan_file = PlanFile.load(self.input().path)
        plan_prompt = plan_file.plan_prompt

        result = ScreenPlanningPrompt.execute(llm, plan_prompt)

        output_raw_path = self.output()['raw'].path
        result.save_raw(output_raw_path)
        output_markdown_path = self.output()['markdown'].path
        result.save_markdown(output_markdown_path)
