#!/usr/bin/env python3
"""
TaskFlow Graph Flow Reproducer
Creates a flow with 52 tasks in waves of 2, where each wave depends on the previous wave.
"""

import logging
import sys
from taskflow import engines
from taskflow.patterns import graph_flow
from taskflow import task

# Configure logging with TRACE level for taskflow
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)

# Enable TRACE level for taskflow (TRACE = 5, below DEBUG)
TRACE_LEVEL = 5
logging.addLevelName(TRACE_LEVEL, 'TRACE')
# Set to 'TRACE' for mor info
LOG_LEVEL='INFO'

# Set taskflow loggers to TRACE level
logging.getLogger('taskflow').setLevel(LOG_LEVEL)

LOG = logging.getLogger(__name__)


class WaveTask(task.Task):
    """A simple task that logs its execution."""

    def __init__(self, name, task_number):
        super(WaveTask, self).__init__(name=name)
        self.task_number = task_number

    def execute(self):
        LOG.info(f"Executing task {self.task_number}: {self.name}")
        return f"Result from task {self.task_number}"


def create_wave_flow():
    """
    Create a graph flow with 52 tasks in waves of 2.

    Wave pattern:
    - Wave 1: Task 1, Task 2 (no dependencies)
    - Wave 2: Task 3, Task 4 (depend on Task 1, Task 2)
    - Wave 3: Task 5, Task 6 (depend on Task 3, Task 4)
    - ...
    - Wave 26: Task 51, Task 52 (depend on Task 49, Task 50)
    """
    flow = graph_flow.Flow('wave_flow')

    # Create all 52 tasks
    tasks = []
    for i in range(1, 53):
        t = WaveTask(name=f'task_{i}', task_number=i)
        tasks.append(t)
        flow.add(t)

    # Add dependencies: each pair depends on the previous pair
    # Tasks are 0-indexed in the list but numbered 1-52
    for i in range(2, 52, 2):
        # Current wave: tasks[i] and tasks[i+1]
        # Previous wave: tasks[i-2] and tasks[i-1]

        # task i+1 depends on task i-1 and task i
        flow.link(tasks[i-2], tasks[i])      # task i depends on task i-1
        flow.link(tasks[i-1], tasks[i])      # task i depends on task i
        flow.link(tasks[i-2], tasks[i+1])    # task i+1 depends on task i-1
        flow.link(tasks[i-1], tasks[i+1])    # task i+1 depends on task i

    return flow


def main():
    """Main execution function."""
    LOG.info("Creating wave flow with 52 tasks...")
    flow = create_wave_flow()

    LOG.info(f"Flow created with {len(list(flow))} tasks")

    LOG.info("Starting flow execution...")
    engine = engines.load(flow, engine='parallel')

    try:
        engine.run()
        LOG.info("Flow execution completed successfully!")
    except Exception as e:
        LOG.error(f"Flow execution failed: {e}")
        raise

    return 0


if __name__ == '__main__':
    sys.exit(main())
