package main

import (
  "context"
  "net/http"
  "time"

  "github.com/cloudwego/hertz/pkg/app"
  "github.com/cloudwego/hertz/pkg/app/server"
  "github.com/cloudwego/hertz/pkg/common/hlog"

  "github.com/hertz-contrib/sse"
)

func main() {
  h := server.Default()

  h.GET("/sse", func(ctx context.Context, c *app.RequestContext) {
    // client can tell server last event it received with Last-Event-ID header
    lastEventID := sse.GetLastEventID(c)
    hlog.CtxInfof(ctx, "last event ID: %s", lastEventID)

    // you must set status code and response headers before first render call
    c.SetStatusCode(http.StatusOK)
    s := sse.NewStream(c)

    count := 0
    sendCountLimit := 10
    for t := range time.NewTicker(1 * time.Second).C {
      event := &sse.Event{
        Event: "timestamp",
        Data:  []byte(t.Format(time.RFC3339)),
      }
      err := s.Publish(event)
      if err != nil {
        return
      }
      count++
      if count >= sendCountLimit {
        // send end flag to client
        err := s.Publish(&sse.Event{
          Event: "end",
          Data:  []byte("end flag"),
        })
        if err != nil {
          return
        }
        break
      }
    }
  })

  h.Spin()
}

package main

import (
  "context"
  "sync"
  "time"

  "github.com/cloudwego/hertz/pkg/common/hlog"

  "github.com/hertz-contrib/sse"
)

var wg sync.WaitGroup

func main() {
  wg.Add(2)
  go func() {
    c := sse.NewClient("http://127.0.0.1:8888/sse")

    // touch off when connected to the server
    c.SetOnConnectCallback(func(ctx context.Context, client *sse.Client) {
      hlog.Infof("client1 connect to server %s success with %s method", c.GetURL(), c.GetMethod())
    })

    // touch off when the connection is shutdown
    c.SetDisconnectCallback(func(ctx context.Context, client *sse.Client) {
      hlog.Infof("client1 disconnect to server %s success with %s method", c.GetURL(), c.GetMethod())
    })

    events := make(chan *sse.Event)
    errChan := make(chan error)
    ctx, cancel := context.WithCancel(context.Background())
    go func() {
      cErr := c.SubscribeWithContext(ctx, func(msg *sse.Event) {
        if msg.Data != nil {
          events <- msg
          return
        }
      })
      errChan <- cErr
    }()
    go func() {
      time.Sleep(5 * time.Second)
      cancel()
      hlog.Info("client1 subscribe cancel")
    }()
    for {
      select {
      case e := <-events:
        hlog.Infof("client1, %+v", e)
      case err := <-errChan:
        if err == nil {
          hlog.Info("client1, ctx done, read stop")
        } else {
          hlog.CtxErrorf(ctx, "client1, err = %s", err.Error())
        }
        wg.Done()
        return
      }
    }
  }()

  go func() {
    c := sse.NewClient("http://127.0.0.1:8888/sse")

    // touch off when connected to the server
    c.SetOnConnectCallback(func(ctx context.Context, client *sse.Client) {
      hlog.Infof("client2 %s connect to server success with %s method", c.GetURL(), c.GetMethod())
    })

    // touch off when the connection is shutdown
    c.SetDisconnectCallback(func(ctx context.Context, client *sse.Client) {
      hlog.Infof("client2 %s disconnect to server success with %s method", c.GetURL(), c.GetMethod())
    })

    events := make(chan *sse.Event, 10)
    errChan := make(chan error)
    go func() {
      cErr := c.Subscribe(func(msg *sse.Event) {
        if msg.Data != nil {
          events <- msg
          return
        }
      })
      errChan <- cErr
    }()

    streamClosed := false
    for {
      select {
      case e := <-events:
        hlog.Infof("client2, %+v", e)
        time.Sleep(2 * time.Second) // do something blocked
        // When the event ends, you should break out of the loop.
        if checkEventEnd(e) {
          wg.Done()
          return
        }
      case err := <-errChan:
        if err == nil {
          // err is nil means read io.EOF, stream is closed
          streamClosed = true
          hlog.Info("client2, stream closed")
          // continue read channel events
          continue
        }
        hlog.CtxErrorf(context.Background(), "client2, err = %s", err.Error())
        wg.Done()
        return
      default:
        if streamClosed {
          hlog.Info("client2, events is empty and stream closed")
          wg.Done()
          return
        }
      }
    }
  }()

  wg.Wait()
}

func checkEventEnd(e *sse.Event) bool {
  // check e.Data or e.Event. It depends on the definition of the server
  return e.Event == "end" || string(e.Data) == "end flag"
}