from crewai import Crew, Process

# IMPORT ALL 5 AGENTS
from agent_01_strategist import StrategistUnit
from agent_02_screenwriter import ScreenwriterUnit
from agent_03_visual import VisualDirectorUnit
from agent_04_motion import MotionDirectorUnit
from agent_05_post import PostSupervisorUnit

# MODULE: STUDIO WORKFLOW
# Responsibility: Connect the wires. Run the factory.
# Address: Root Folder / studio_workflow.py

def run_production_studio(topic):
    """
    Triggers the FULL 5-Stage Linear Assembly Line:
    Strat -> Writer -> Visual -> Motion -> Post
    """
    
    # 1. Initialize Units
    unit_01 = StrategistUnit()
    unit_02 = ScreenwriterUnit()
    unit_03 = VisualDirectorUnit()
    unit_04 = MotionDirectorUnit()
    unit_05 = PostSupervisorUnit()
    
    # 2. Create Agents (The Workers)
    agent_01 = unit_01.create_agent()
    agent_02 = unit_02.create_agent()
    agent_03 = unit_03.create_agent()
    agent_04 = unit_04.create_agent()
    agent_05 = unit_05.create_agent()
    
    # 3. Create Tasks (The Jobs)
    # The output of Task 1 automatically feeds into Task 2 because of the order below
    task_01 = unit_01.create_task(topic)
    task_02 = unit_02.create_task(task_01) 
    task_03 = unit_03.create_task(task_02)
    task_04 = unit_04.create_task(task_03)
    task_05 = unit_05.create_task(task_04)
    
    # 4. Assemble the Crew
    studio_crew = Crew(
        agents=[agent_01, agent_02, agent_03, agent_04, agent_05],
        tasks=[task_01, task_02, task_03, task_04, task_05],
        process=Process.sequential, # Linear assembly line
        verbose=True
    )
    
    # 5. Action!
    # Kickoff returns the result of the LAST task (The Post-Production Plan)
    result = studio_crew.kickoff()
    return result

if __name__ == "__main__":
    print("Testing Full Studio Workflow...")
    # WARNING: This will take a while on local hardware!
    print(run_production_studio("The History of Coffee in 30 seconds"))