import typing as ty
from collections import OrderedDict
from dataclasses import dataclass

import pytest
from dataclasses_json import dataclass_json
from flytekit import PythonInstanceTask
from flytekit import task
from flytekit import workflow
from flytekit.configuration import Image
from flytekit.configuration import ImageConfig
from flytekit.configuration import SerializationSettings
from flytekit.core.base_task import PythonTask
from flytekit.core.interface import Interface
from flytekit.tools.translator import get_serializable
from flytekit.types.file import FlyteFile


@dataclass_json
@dataclass
class MyTaskConfig:
    string: str


@dataclass_json
@dataclass
class RuntimeOutput:
    model: FlyteFile["tar.gz"]


class MyTask(PythonTask[MyTaskConfig]):
    _TASK_TYPE = "unittest_my_task"

    def __init__(
        self,
        name: str,
        task_config: MyTaskConfig,
        inputs: ty.Dict[str, ty.Any],
        **kwargs,
    ):
        if not isinstance(task_config, MyTaskConfig):
            raise TypeError("Invalid task_config type")

        outputs = {"result": RuntimeOutput}
        super().__init__(
            name=name,
            task_config=task_config,
            task_type=self._TASK_TYPE,
            interface=Interface(inputs=inputs, outputs=outputs),
            **kwargs,
        )

    def execute(self, **kwargs) -> RuntimeOutput:
        print(kwargs)
        return RuntimeOutput(model=FlyteFile("<s3://bucket/model.tar.gz>"))


def test_workflow_serializes():
    @task
    def get_data() -> ty.Tuple[FlyteFile, FlyteFile]:
        return FlyteFile("<s3://bucket/train.csv>"), FlyteFile(
            "<s3://s3_bucket/valid.csv>"
        )

    task_config = MyTaskConfig(
        string="some string",
    )
    training_task = MyTask(
        "unittest-mytask",
        task_config,
        inputs={"train": FlyteFile["csv"], "valid": FlyteFile["csv"]},
    )

    @workflow
    def test_ml_pipeline() -> RuntimeOutput:
        train, valid = get_data()
        training_output = training_task(train=train, valid=valid)
        return training_output

    default_img = Image(name="default", fqn="test", tag="tag")
    serialization_settings = SerializationSettings(
        project="project",
        domain="domain",
        version="version",
        env=None,
        image_config=ImageConfig(default_image=default_img, images=[default_img]),
    )

    wf_spec = get_serializable(OrderedDict(), serialization_settings, test_ml_pipeline)
    assert wf_spec is not None
    assert wf_spec.template is not None

logger.debug(f"Looking for LHS for {self} from {self._instantiated_in}")
        m = _importlib.import_module(self._instantiated_in)
        for k in dir(m):
            try:
                if getattr(m, k) is self:
                    logger.debug(f"Found LHS for {self}, {k}")
                    self._lhs = k
                    return k
            except ValueError as err:
                # Empty pandas dataframes behave weirdly here such that calling `m.df` raises:
                # ValueError: The truth value of a {type(self).__name__} is ambiguous. Use a.empty, a.bool(), a.item(),
                #   a.any() or a.all()
                # Since dataframes aren't registrable entities to begin with we swallow any errors they raise and
                # continue looping through m.
                logger.warning("Caught ValueError {} while attempting to auto-assign name".format(err))
                pass
    
        logger.error(f"Could not find LHS for {self} in {self._instantiated_in}")
>       raise _system_exceptions.FlyteSystemException(f"Error looking for LHS in {self._instantiated_in}")
E       flytekit.exceptions.system.FlyteSystemException: Error looking for LHS in __main__

venv/lib/python3.8/site-packages/flytekit/core/tracker.py:96: FlyteSystemException