package websocket

import (
	"context"
	"sync"
	"sync/atomic"
	"time"

	"github.com/gogf/gf/v2/encoding/gjson"
	"github.com/gogf/gf/v2/os/glog"
	"github.com/gorilla/websocket"
)

// WsConnection 单个WebSocket连接，Metadata 存放业务自定义数据
type WsConnection struct {
	SessionId string
	Conn      *websocket.Conn
	Headers   map[string]string
	Metadata  sync.Map // 业务数据：FlowId, execCancel 等

	writeMu     sync.Mutex // 保护 websocket.Conn 并发写（WriteMessage + WriteControl）
	closeCancel context.CancelFunc
	closed      int32
}

// SetMeta 设置业务元数据
func (c *WsConnection) SetMeta(key string, value interface{}) {
	c.Metadata.Store(key, value)
}

// GetMeta 获取业务元数据
func (c *WsConnection) GetMeta(key string) (interface{}, bool) {
	return c.Metadata.Load(key)
}

// GetMetaT 泛型版 GetMeta，省去外部类型断言
func GetMetaT[T any](c *WsConnection, key string) (T, bool) {
	val, ok := c.Metadata.Load(key)
	if !ok {
		var zero T
		return zero, false
	}
	t, ok := val.(T)
	return t, ok
}

// IsClosed 连接是否已关闭
func (c *WsConnection) IsClosed() bool {
	return atomic.LoadInt32(&c.closed) == 1
}

// WriteControl 带写锁保护的 WriteControl，用于心跳 Ping / Pong / Close帧
func (c *WsConnection) WriteControl(msgType int, data []byte, deadline time.Time) error {
	c.writeMu.Lock()
	defer c.writeMu.Unlock()
	_ = c.Conn.SetWriteDeadline(deadline)
	return c.Conn.WriteControl(msgType, data, deadline)
}

// WriteJSON 业务层外部写入入口，共享 writeMu 与心跳/Pong 互斥
func (c *WsConnection) WriteJSON(data interface{}) error {
	jsonBytes, err := gjson.Encode(data)
	if err != nil {
		glog.Errorf(context.Background(), "json encode failed: %v", err)
		return err
	}
	c.writeMu.Lock()
	_ = c.Conn.SetWriteDeadline(time.Now().Add(30 * time.Second))
	err = c.Conn.WriteMessage(websocket.TextMessage, jsonBytes)
	c.writeMu.Unlock()
	if err != nil {
		glog.Debugf(context.Background(), "websocket write failed: %v", err)
	}
	return err
}