package main

import (
	"context"
	"log"
	"os"
	"time"

	amqp "github.com/rabbitmq/amqp091-go"
)

// 连接字符串
const amqpURL = "amqp://guest:guest@localhost:5672/"

// 队列名称
const queueName = "hello"

// 错误处理函数
func failOnError(err error, msg string) {
	if err != nil {
		log.Panicf("%s: %s", msg, err)
	}
}

//---

// 生产者函数
func producer() {
	// 1. 连接 RabbitMQ 服务器
	conn, err := amqp.Dial(amqpURL)
	failOnError(err, "Failed to connect to RabbitMQ")
	defer conn.Close()

	// 2. 创建通道
	ch, err := conn.Channel()
	failOnError(err, "Failed to open a channel")
	defer ch.Close()

	// 3. 声明队列（如果队列不存在则创建）
	q, err := ch.QueueDeclare(
		queueName, // name
		false,     // durable 持久化
		false,     // delete when unused 当没有消费者时删除队列
		false,     // exclusive 独占
		false,     // no-wait
		nil,       // arguments
	)
	failOnError(err, "Failed to declare a queue")

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	// 4. 发布消息
	body := "Hello RabbitMQ!"
	err = ch.PublishWithContext(ctx,
		"",     // exchange
		q.Name, // routing key
		false,  // mandatory
		false,  // immediate
		amqp.Publishing{
			ContentType: "text/plain",
			Body:        []byte(body),
		})
	failOnError(err, "Failed to publish a message")
	log.Printf(" [x] Sent %s\n", body)
}

//---

// 消费者函数
func consumer() {
	// 1. 连接 RabbitMQ 服务器
	conn, err := amqp.Dial(amqpURL)
	failOnError(err, "Failed to connect to RabbitMQ")
	defer conn.Close()

	// 2. 创建通道
	ch, err := conn.Channel()
	failOnError(err, "Failed to open a channel")
	defer ch.Close()

	// 3. 声明队列（消费者也需要确保队列存在）
	q, err := ch.QueueDeclare(
		queueName, // name
		false,     // durable
		false,     // delete when unused
		false,     // exclusive
		false,     // no-wait
		nil,       // arguments
	)
	failOnError(err, "Failed to declare a queue")

	// 4. 注册消费者
	msgs, err := ch.Consume(
		q.Name, // queue
		"",     // consumer
		true,   // auto-ack 自动确认
		false,  // exclusive
		false,  // no-local
		false,  // no-wait
		nil,    // args
	)
	failOnError(err, "Failed to register a consumer")

	var forever chan struct{}

	// 5. 持续消费消息
	go func() {
		for d := range msgs {
			log.Printf(" [x] Received a message: %s", d.Body)
		}
	}()

	log.Printf(" [*] Waiting for messages. To exit press CTRL+C")
	<-forever
}

//---

// 主函数，根据命令行参数启动生产者或消费者
func main() {
	if len(os.Args) < 2 {
		log.Fatalf("Usage: go run main.go [producer|consumer]")
	}

	mode := os.Args[1]
	switch mode {
	case "producer":
		log.Println("Starting producer...")
		producer()
	case "consumer":
		log.Println("Starting consumer...")
		consumer()
	default:
		log.Fatalf("Invalid mode: %s. Use 'producer' or 'consumer'.", mode)
	}
}