package main

import (
	"log"
	"net/url"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

func main() {
	router := gin.Default()
	router.GET("/connect", func(c *gin.Context) {
		connectWebSocket()
	})
	router.Run(":8080") // 监听并在 0.0.0.0:8080 上启动服务
}

func connectWebSocket() {
	// WebSocket 服务地址
	u := url.URL{Scheme: "ws", Host: "58.34.54.45:13999", Path: "/queue/join"}
	log.Printf("Connecting to %s", u.String())

	// 连接 WebSocket 服务器
	conn, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
	if err != nil {
		log.Fatal("dial:", err)
	}
	defer conn.Close()

	done := make(chan struct{})

	// 接收消息
	go func() {
		defer close(done)
		for {
			_, message, err := conn.ReadMessage()
			if err != nil {
				log.Println("read:", err)
				return
			}
			log.Printf("recv: %s", message)

			// 根据接收的消息进行处理
			if string(message) == `{"msg":"send_data"}` {
				// 当服务器请求数据时，发送音频文件信息
				sendAudioFileInfo(conn)
			}
		}
	}()

	// 发送初始参数
	initialParams := `{"fn_index":0,"session_hash":"7rc3vrzkgf4"}`
	err = conn.WriteMessage(websocket.TextMessage, []byte(initialParams))
	if err != nil {
		log.Println("write:", err)
		return
	}

	// 等待一段时间，让消息传输完成
	time.Sleep(1 * time.Second)

	// 关闭连接
	err = conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
	if err != nil {
		log.Println("write close:", err)
		return
	}
	<-done
}

func sendAudioFileInfo(conn *websocket.Conn) {
	// 音频文件信息
	audioInfo := `{"data":[[{"data":"http://58.34.54.45:13999/file=/tmp/gradio/899ae5805ee3016b2bfad387345b285eee29c5ad/test.mp3","name":"/tmp/gradio/899ae5805ee3016b2bfad387345b285eee29c5ad/test.mp3","size":423208,"orig_name":"test.mp3","is_file":true}],"medium","Automatic Detection","txt",false,true,1,-1,0.6,"float16"],"event_data":null,"fn_index":0,"session_hash":"7rc3vrzkgf4"}`
	err := conn.WriteMessage(websocket.TextMessage, []byte(audioInfo))
	if err != nil {
		log.Println("write:", err)
		return
	}
}
