import time
import dataclasses

from pyflink.common import Types, Row
from pyflink.common import WatermarkStrategy
from pyflink.common.watermark_strategy import TimestampAssigner
from pyflink.datastream import StreamExecutionEnvironment
from pyflink.datastream.window import TumblingEventTimeWindows, Time


@dataclasses.dataclass
class MyData:
    value: int
    key: str
    marker: str

    @classmethod
    def from_row(cls, row: Row):
        return cls(value=row.value, key=row.key, marker=row.marker)

    @staticmethod
    def to_stats(row: Row):
        data = MyData.from_row(row)
        return Stats(data.value, data.key)

    @staticmethod
    def get_type_info():
        return Types.ROW_NAMED(
            field_names=[
                "value",
                "key",
                "marker",
            ],
            field_types=[
                Types.INT(),
                Types.STRING(),
                Types.STRING(),
            ],
        )


@dataclasses.dataclass
class BaseStats:
    value: int
    key: str


@dataclasses.dataclass
class Stats(BaseStats):
    @classmethod
    def from_mydata(cls, data: MyData):
        return cls(value=data.value, key=data.key)

    @staticmethod
    def merge(this: BaseStats, that: BaseStats):
        return Stats(value=this.value + that.value, key=this.key)


class CustomTimestampAssigner(TimestampAssigner):
    def extract_timestamp(self, value, record_timestamp):
        if value.marker == "late":
            return int(time.time_ns() / 1000000) + 50000
        else:
            return int(time.time_ns() / 1000000)


custom_watermark_strategy = WatermarkStrategy.for_monotonous_timestamps().with_timestamp_assigner(
    CustomTimestampAssigner()
)

env = StreamExecutionEnvironment.get_execution_environment()

data_stream = env.from_collection(
    collection=[(1, "a", "on_time"), (2, "a", "on_time"), (3, "a", "late")],
    type_info=MyData.get_type_info(),
).assign_timestamps_and_watermarks(custom_watermark_strategy)

data_stream.map(MyData.from_row).map(Stats.from_mydata).key_by(lambda s: s.key).window(
    TumblingEventTimeWindows.of(Time.minutes(1))
).reduce(Stats.merge).print()

env.execute()