from collections import namedtuple

import bonobo
from bonobo.config import use_raw_input
from bonobo.execution.contexts import GraphExecutionContext
from bonobo.util.bags import BagType

Extracted = namedtuple('Extracted', ['id', 'name', 'value'])
ExtractedBT = BagType('ExtractedBT', ['id', 'name', 'value'])


def extract_nt():
    yield Extracted(id=1, name='Guido', value='.py')
    yield Extracted(id=2, name='Larry', value='.pl')
    yield Extracted(id=3, name='Dennis', value='.c')
    yield Extracted(id=4, name='Yukihiro', value='.rb')


def extract_bt():
    yield ExtractedBT(id=1, name='Guido', value='.py')
    yield ExtractedBT(id=2, name='Larry', value='.pl')
    yield ExtractedBT(id=3, name='Dennis', value='.c')
    yield ExtractedBT(id=4, name='Yukihiro', value='.rb')


def transform_using_args(id, name, value):
    yield Extracted(id=id * 2, name=name, value=name.lower() + value)


@use_raw_input
def transform_nt(row):
    yield row._replace(name=row.name.upper())


def StoreInList(buffer: list):
    def store_in_list(*args, buffer=buffer):
        buffer.append(args)

    return store_in_list


def test_execution():
    graph = bonobo.Graph()

    result_args = []
    result_nt = []
    result_bt = []

    graph.add_chain(extract_nt, transform_using_args, StoreInList(result_args))
    graph.add_chain(transform_nt, StoreInList(result_nt), _input=extract_nt)
    graph.add_chain(extract_bt, transform_using_args, StoreInList(result_bt))

    with GraphExecutionContext(graph) as context:
        context.run_until_complete()

    assert result_args == [(2, 'Guido', 'guido.py'), (4, 'Larry', 'larry.pl'), (6, 'Dennis', 'dennis.c'),
                           (8, 'Yukihiro', 'yukihiro.rb')]

    assert result_nt == [(1, 'GUIDO', '.py'), (2, 'LARRY', '.pl'), (3, 'DENNIS', '.c'), (4, 'YUKIHIRO', '.rb')]

    assert result_bt == [(2, 'Guido', 'guido.py'), (4, 'Larry', 'larry.pl'), (6, 'Dennis', 'dennis.c'),
                         (8, 'Yukihiro', 'yukihiro.rb')]