from __future__ import annotations

from dataclasses import dataclass

from flytekit import task, workflow
from mashumaro.mixins.json import DataClassJSONMixin


@dataclass
class Fruit(DataClassJSONMixin):
    name: str


@dataclass
class NestedFruit(DataClassJSONMixin):
    sub_fruit: Fruit
    name: str


@task
def dataclass_task() -> Fruit:
    return Fruit(name="banana")


@task
def dataclass_task_nested() -> NestedFruit:
    return NestedFruit(sub_fruit=Fruit(name="banana"), name="nested_name")


@task
def print_message(message: str):
    print(message)
    return


@task
def print_message_fruit(fruit_instance: Fruit):
    print(fruit_instance.name)
    print(fruit_instance)
    return


@task
def print_message_nested(nested_fruit: NestedFruit):
    print(nested_fruit.sub_fruit.name)
    print(nested_fruit.name)
    return


@workflow
def dataclass_wf():
    fruit_instance = dataclass_task()
    nested_fruit_instance = dataclass_task_nested()

    # ✅ works
    print_message(message=fruit_instance.name)
    # ✅ works
    print_message_nested(nested_fruit=nested_fruit_instance)
    # ❌ fails
    print_message_fruit(fruit_instance=nested_fruit_instance.sub_fruit)

Error 0: Code: MismatchingTypes, Node Id: n4, Description: Variable [o0] (type [simple:STRING]) doesn't match expected type [simple:STRUCT  metadata:{fields:{key:"additionalProperties"  value:{bool_value:false}}  fields:{key:"properties"  value:{struct_value:{fields:{key:"name"  value:{struct_value:{fields:{key:"type"  value:{string_value:"string"}}}}}}}}  fields:{key:"required"  value:{list_value:{values:{string_value:"name"}}}}  fields:{key:"title"  value:{string_value:"Fruit"}}  fields:{key:"type"  value:{string_value:"object"}}}  structure:{dataclass_type:{key:"name"  value:{simple:STRING}}}].
Error 1: Code: ParameterNotBound, Node Id: n4, Description: Parameter not bound [fruit_instance]