package main

import (
    "context"
    "fmt"
    "log"
    "os"
    "os/signal"
    "syscall"
    "time"

    "github.com/segmentio/kafka-go"
)

var (
    topic  = "user_click"
    reader *kafka.Reader
)

func kafkaConn() *kafka.Conn {
    topic := "my-toc"
    partition := 0

    conn, err := kafka.DialLeader(context.Background(), "tcp", "42.193.152.252:9092", topic, partition)
    if err != nil {
        log.Fatal("failed to dial leader:", err)
    }
    // kafkaW(conn)
    return conn

}

func kafkaW(conn *kafka.Conn) {
    conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
    _, err := conn.WriteMessages(
        kafka.Message{Value: []byte("one!")},
        kafka.Message{Value: []byte("two!")},
        kafka.Message{Value: []byte("three!")},
    )
    if err != nil {
        log.Fatal("failed to write messages:", err)
    }

    if err := conn.Close(); err != nil {
        log.Fatal("failed to close writer:", err)
    }
    fmt.Println("写入成功")
}

func kafkaR(conn *kafka.Conn) {
    fmt.Println("开始读取")
    conn.SetReadDeadline(time.Now().Add(10 * time.Second))
    batch := conn.ReadBatch(10e3, 1e6) // fetch 10KB min, 1MB max

    b := make([]byte, 10e3) // 10KB max per message
    for {
        n, err := batch.Read(b)
        if err != nil {
            break
        }
        fmt.Println(string(b[:n]))
    }

    if err := batch.Close(); err != nil {
        log.Fatal("failed to close batch:", err)
    }

    if err := conn.Close(); err != nil {
        log.Fatal("failed to close connection:", err)
    }
    fmt.Println("读取完毕")
}

// 生产消息
func writeKafka(ctx context.Context) {
    writer := &kafka.Writer{
        Addr:                   kafka.TCP("47.120.55.99:9092"), //不定长参数，支持传入多个broker的ip:port
        Topic:                  topic,                           //为所有message指定统一的topic。如果这里不指定统一的Topic，则创建kafka.Message{}时需要分别指定Topic
        Balancer:               &kafka.Hash{},                   //把message的key进行hash，确定partition
        WriteTimeout:           1 * time.Second,                 //设定写超时
        RequiredAcks:           kafka.RequireNone,               //RequireNone不需要等待ack返回，效率最高，安全性最低；RequireOne只需要确保Leader写入成功就可以发送下一条消息；RequiredAcks需要确保Leader和所有Follower都写入成功才可以发送下一条消息。
        AllowAutoTopicCreation: true,                            //Topic不存在时自动创建。生产环境中一般设为false，由运维管理员创建Topic并配置partition数目
    }
    defer writer.Close() //记得关闭连接

    for i := 0; i < 3; i++ { //允许重试3次
        if err := writer.WriteMessages(ctx, //批量写入消息，原子操作，要么全写成功，要么全写失败
            kafka.Message{Key: []byte("1"), Value: []byte("大")},
            kafka.Message{Key: []byte("2"), Value: []byte("乔")},
            kafka.Message{Key: []byte("3"), Value: []byte("乔")},
            kafka.Message{Key: []byte("1"), Value: []byte("教")}, //key相同时肯定写入同一个partition
            kafka.Message{Key: []byte("2"), Value: []byte("育")},
        ); err != nil {
            if err == kafka.LeaderNotAvailable { //首次写一个新的Topic时，会发生LeaderNotAvailable错误，重试一次就好了
                time.Sleep(500 * time.Millisecond)
                continue
            } else {
                fmt.Printf("kafka写入失败: %v", err)
            }
        } else {
            fmt.Println("写入成功")
            break //只要成功一次就不再尝试下一次了
        }
    }
}

// 消费消息
func readKafka(ctx context.Context) {
    reader = kafka.NewReader(kafka.ReaderConfig{
        Brokers:        []string{"47.120.55.99:9092"}, //支持传入多个broker的ip:port
        Topic:          topic,
        CommitInterval: 1 * time.Second,   //每隔多长时间自动commit一次offset。即一边读一边向kafka上报读到了哪个位置。
        GroupID:        "recommend_biz",   //一个Group内消费到的消息不会重复
        StartOffset:    kafka.FirstOffset, //当一个特定的partition没有commited offset时(比如第一次读一个partition，之前没有commit过)，通过StartOffset指定从第一个还是最后一个位置开始消费。StartOffset的取值要么是FirstOffset要么是LastOffset，LastOffset表示Consumer启动之前生成的老数据不管了。仅当指定了GroupID时，StartOffset才生效。
    })
    // defer reader.Close() //由于下面是死循环，正常情况下readKafka()函数永远不会结束，defer不会执行。所以需要监听信息2和15，当收到信号时关闭reader。需要把reader设为全局变量

    for { //消息队列里随时可能有新消息进来，所以这里是死循环，类似于读Channel
        if message, err := reader.ReadMessage(ctx); err != nil {
            fmt.Printf("kafka读取失败: %v", err)
            break
        } else {
            offset := message.Offset
            fmt.Printf("topic=%s, partition=%d, offset=%d, key=%s, message content=%s\n", message.Topic, message.Partition, offset, string(message.Key), string(message.Value))
        }
    }
}

// 需要监听信息2和15，当收到信号时关闭reader, 由于执行kill时无法触发defer关闭连接, 所以需要在主进程进行监听, 关闭信号
func listenSignal() {
    c := make(chan os.Signal, 1)                      // 创建chan管道
    signal.Notify(c, syscall.SIGINT, syscall.SIGTERM) // 注册信号2和15到管道里
    sig := <-c                                        // 阻塞，直到信号的到来
    fmt.Printf("收到信号: %s\n", sig.String())            // ctrl + c
    if reader != nil {
        reader.Close() // 关闭kafka连接
    }
    os.Exit(0) //进程退出
}

func main() {
    ctx := context.Background()
    // writeKafka(ctx)

    go listenSignal()
    readKafka(ctx)
    // kafkaW(kafkaConn())
    // kafkaR(kafkaConn())
}