import sys
from pprint import pprint

from haystack import Pipeline, component


@component
class QuestionClassifier:
    @component.output_types(topic=str)
    def run(self, question: str):
        normalized = " ".join(question.strip().split()).lower()
        if "pipeline" in normalized:
            return {"topic": "pipeline"}
        return {"topic": "general"}


@component
class ReplyBuilder:
    @component.output_types(reply=str)
    def run(self, topic: str):
        return {"reply": f"route={topic}"}


pipeline = Pipeline()
pipeline.add_component("classifier", QuestionClassifier())
pipeline.add_component("reply", ReplyBuilder())
pipeline.connect("classifier.topic", "reply.topic")

question = " ".join(sys.argv[1:]) or "How do Haystack pipelines run?"
data = {"classifier": {"question": question}}

print("input:")
pprint(data)

print("\ninputs:")
pprint(pipeline.inputs())

print("\nleaf output:")
pprint(pipeline.run(data))

print("\ninclude classifier output:")
pprint(pipeline.run(data, include_outputs_from={"classifier"}))