import dapr.ext.workflow as wf
from dotenv import load_dotenv
from openai import OpenAI
from time import sleep

# Load environment variables
load_dotenv()

# Initialize Workflow Instance
wfr = wf.WorkflowRuntime()


# Define Workflow logic
@wfr.workflow(name="task_chain_workflow")
def task_chain_workflow(ctx: wf.DaprWorkflowContext):
    result1 = yield ctx.call_activity(get_character)
    result2 = yield ctx.call_activity(get_line, input=result1)
    return result2


# Activity 1
@wfr.activity(name="step1")
def get_character(ctx):
    client = OpenAI()
    response = client.chat.completions.create(
        messages=[
            {
                "role": "user",
                "content": "Pick a random character from The Lord of the Rings and respond with the character name only",
            }
        ],
        model="gpt-4o",
    )
    character = response.choices[0].message.content
    print(f"Character: {character}")
    return character


# Activity 2
@wfr.activity(name="step2")
def get_line(ctx, character: str):
    client = OpenAI()
    response = client.chat.completions.create(
        messages=[{"role": "user", "content": f"What is a famous line by {character}"}],
        model="gpt-4o",
    )
    line = response.choices[0].message.content
    print(f"Line: {line}")
    return line


if __name__ == "__main__":
    wfr.start()
    sleep(5)  # wait for workflow runtime to start

    wf_client = wf.DaprWorkflowClient()
    instance_id = wf_client.schedule_new_workflow(workflow=task_chain_workflow)
    print(f"Workflow started. Instance ID: {instance_id}")
    state = wf_client.wait_for_workflow_completion(instance_id)
    print(f"Workflow completed! Status: {state.runtime_status}")

    wfr.shutdown()