import logging
import apache_beam as beam
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.coders import coders
from apache_beam.transforms.userstate import BagStateSpec
from apache_beam.transforms.userstate import TimerSpec
from apache_beam.transforms.userstate import on_timer
from apache_beam.transforms.timeutil import TimeDomain
from apache_beam.transforms.periodicsequence import PeriodicImpulse

logging.basicConfig(level=logging.INFO)

options = PipelineOptions([
    "--streaming",
    "--environment_type=LOOPBACK",
    "--runner=PrismRunner",
    "--prism_beam_version_override=v2.65.0"
    #"--runner=PortableRunner", "--job_endpoint=localhost:8073",
])

class MyStatefulDoFn(beam.DoFn):
  ALL_ELEMENTS = BagStateSpec('buffer', coders.FloatCoder())
  TIMER = TimerSpec('timer', TimeDomain.WATERMARK)

  def __init__(self, n):
    self._n = n

  def process(self, element_pair,
              t = beam.DoFn.TimestampParam,
              state=beam.DoFn.StateParam(ALL_ELEMENTS),
              timer = beam.DoFn.TimerParam(TIMER)
              ):
    state.add(element_pair[1])
    all_elements = list(state.read())

    # the start of new state cell
    if len(all_elements) == 1:
      print("set timer to", t+10)
      timer.set(t+10)

    # the state cell is full
    if len(all_elements) == self._n:
      print("batch process start")
      for e in all_elements:
        yield e
      print("batch process end")

      # clear the buffer
      state.clear()

      # don't need the timer now
      print("clear timer")
      timer.clear()

  @on_timer(TIMER)
  def timer_callback(self, t = beam.DoFn.TimestampParam, state = beam.DoFn.StateParam(ALL_ELEMENTS)):
    print("TIMER SHOULD NOT FIRE HERE!")
    print(f"timer callback start (timestamp={t})")
    all_elements = list(state.read())
    print("timer callback elements:", all_elements)
    print("timer callback end")


with beam.Pipeline(options=options) as p:
  _ = (
      p | PeriodicImpulse(fire_interval=1)
      | beam.WithKeys(0)
      | beam.ParDo(MyStatefulDoFn(5))
      | beam.Map(print))