package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"os"
	"time"

	"github.com/redis/go-redis/v9"
)

// 定义 Redis 配置和频道名称
const redisAddr = "localhost:6379" // 替换成您的 Redis 地址
const channelName = "stream:10"

// 全局 Redis 客户端
var rdb *redis.Client

// 定义命令行参数
var mode = flag.String("mode", "", "运行模式：必须是 'pub' (发布者) 或 'sub' (订阅者)")

// 初始化 Redis 客户端并检查连接
func initRedis() {
	rdb = redis.NewClient(&redis.Options{
		Addr:     redisAddr,
		Password: "", // 您的密码（如果没有则为空）
		DB:       7,  // 默认 DB
	})

	ctx := context.Background()
	_, err := rdb.Ping(ctx).Result()
	if err != nil {
		log.Fatalf("无法连接到 Redis (%s): %v", redisAddr, err)
	}
	fmt.Println("成功连接到 Redis。")
}

// 订阅者逻辑
func runSubscriber() {
	fmt.Println("--- 启动订阅者 (Subscriber) 模式 ---")

	// 订阅频道
	// Subscribes to the given channels.
	pubsub := rdb.Subscribe(context.Background(), channelName)
	defer pubsub.Close() // 确保在函数退出时关闭连接

	// 接收并处理消息
	fmt.Printf("已订阅频道: %s。等待消息...\n", channelName)

	// 使用 Channel() 方法获取一个 Go Channel 来接收消息
	ch := pubsub.Channel()

	// 循环接收消息
	for msg := range ch {
		fmt.Printf("\n--- 收到消息 ---\n")
		fmt.Printf("频道 (Channel): %s\n", msg.Channel)
		fmt.Printf("消息 (Payload): %s\n", msg.Payload)
		fmt.Printf("----------------\n")

		// 添加退出逻辑
		if msg.Payload == "exit" {
			fmt.Println("\n收到 'exit' 消息，退出订阅循环。")
			break
		}
	}

	fmt.Println("订阅者程序退出。")
}

// 发布者逻辑
func runPublisher() {
	fmt.Println("--- 启动发布者 (Publisher) 模式 ---")

	messages := []string{
		"Hello Redis Pub/Sub!",
		"This is message two from publisher.",
		"A third message is here.",
		"exit", // 发送退出信号给订阅者
	}

	ctx := context.Background()

	// 循环发布消息
	for i, msg := range messages {
		// Publish publishes message to channel.
		// Result() 返回值是接收到消息的订阅者数量 (int64)
		result, err := rdb.Publish(ctx, channelName, msg).Result()
		if err != nil {
			log.Fatalf("发布消息失败: %v", err)
		}

		fmt.Printf("发布消息 %d: '%s' 到频道 '%s'，有 %d 个订阅者接收。\n",
			i+1, msg, channelName, result)

		// 每次发送后等待一小段时间
		time.Sleep(500 * time.Millisecond)
	}

	fmt.Println("发布者程序完成。")
}

func main() {
	// 解析命令行参数
	flag.Parse()

	// 检查模式参数
	if *mode != "pub" && *mode != "sub" {
		fmt.Println("错误：请指定有效的运行模式。")
		flag.Usage()
		os.Exit(1)
	}

	// 初始化 Redis
	initRedis()

	// 根据模式执行相应逻辑
	if *mode == "sub" {
		runSubscriber()
	} else if *mode == "pub" {
		runPublisher()
	}
}
