package minecraftquery

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"strings"
	"sync"
	"time"
)

// 连接池配置
var (
	connPool   sync.Map           // 存储活跃连接
	poolMutex  sync.Mutex         // 连接池互斥锁
	connExpiry = 30 * time.Second // 连接最长保留时间
)

type ServerStatus struct {
	Version      string
	Protocol     int
	Online       int
	MaxPlayers   int
	Description  string
	Latency      time.Duration
	ResolvedHost string // 实际连接的服务器地址
}

func Query(addr string, port int, timeout time.Duration) (*ServerStatus, error) {
	start := time.Now()

	// 解析SRV记录
	resolvedAddr, resolvedPort, err := resolveSRV(addr)
	if err == nil {
		addr, port = resolvedAddr, resolvedPort
	}

	key := fmt.Sprintf("%s:%d", addr, port)

	// 尝试获取连接
	conn, err := getConnection(key, timeout)
	if err != nil {
		return nil, err
	}
	defer returnConnection(key, conn)

	// 执行查询
	status, err := performQuery(conn, addr, port, start)
	if err != nil {
		conn.Close()
		connPool.Delete(key)
		return nil, err
	}

	status.ResolvedHost = fmt.Sprintf("%s:%d", addr, port)
	return status, nil
}

// SRV记录解析
func resolveSRV(domain string) (string, int, error) {
	_, addrs, err := net.LookupSRV("minecraft", "tcp", domain)
	if err != nil || len(addrs) == 0 {
		return "", 0, errors.New("no SRV record found")
	}

	// 选择优先级最高且权重最大的记录
	best := addrs[0]
	for _, a := range addrs[1:] {
		if a.Priority < best.Priority ||
			(a.Priority == best.Priority && a.Weight > best.Weight) {
			best = a
		}
	}

	return strings.TrimSuffix(best.Target, "."), int(best.Port), nil
}

// 连接管理
func getConnection(key string, timeout time.Duration) (net.Conn, error) {
	// 尝试获取现有连接
	if val, ok := connPool.Load(key); ok {
		conn := val.(*pooledConn)
		if time.Since(conn.lastUsed) < connExpiry && isConnAlive(conn.Conn) {
			poolMutex.Lock()
			conn.lastUsed = time.Now()
			poolMutex.Unlock()
			return conn.Conn, nil
		}
		conn.Close()
		connPool.Delete(key)
	}

	// 创建新连接
	conn, err := net.DialTimeout("tcp", key, timeout)
	if err != nil {
		return nil, fmt.Errorf("connection failed: %w", err)
	}

	pooled := &pooledConn{
		Conn:     conn,
		lastUsed: time.Now(),
	}
	connPool.Store(key, pooled)
	return conn, nil
}

func returnConnection(key string, conn net.Conn) {
	if val, ok := connPool.Load(key); ok {
		pooled := val.(*pooledConn)
		poolMutex.Lock()
		pooled.lastUsed = time.Now()
		poolMutex.Unlock()
	}
}

type pooledConn struct {
	net.Conn
	lastUsed time.Time
}

func isConnAlive(conn net.Conn) bool {
	// 发送1字节的无效数据测试连接
	conn.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
	defer conn.SetReadDeadline(time.Time{})

	_, err := conn.Write([]byte{0x00})
	if err != nil {
		return false
	}

	buf := make([]byte, 1)
	_, err = conn.Read(buf)
	return err == nil || errors.Is(err, io.EOF)
}

// 查询实现
func performQuery(conn net.Conn, addr string, port int, start time.Time) (*ServerStatus, error) {
	conn.SetDeadline(time.Now().Add(5 * time.Second))

	if err := sendHandshake(conn, addr, port); err != nil {
		return nil, fmt.Errorf("handshake failed: %w", err)
	}

	if err := sendStatusRequest(conn); err != nil {
		return nil, fmt.Errorf("status request failed: %w", err)
	}

	jsonData, err := readStatusResponse(conn)
	if err != nil {
		return nil, fmt.Errorf("response parse failed: %w", err)
	}

	var response struct {
		Version struct {
			Name     string `json:"name"`
			Protocol int    `json:"protocol"`
		} `json:"version"`
		Players struct {
			Online int `json:"online"`
			Max    int `json:"max"`
		} `json:"players"`
		Description json.RawMessage `json:"description"`
	}
	if err := json.Unmarshal(jsonData, &response); err != nil {
		return nil, fmt.Errorf("json parse error: %w", err)
	}

	return &ServerStatus{
		Version:     response.Version.Name,
		Protocol:    response.Version.Protocol,
		Online:      response.Players.Online,
		MaxPlayers:  response.Players.Max,
		Description: parseDescription(response.Description),
		Latency:     time.Since(start),
	}, nil
}

// 以下为内部实现细节
func writeVarInt(w io.Writer, value int) error {
	for {
		if (value & ^0x7F) == 0 {
			return binary.Write(w, binary.BigEndian, byte(value))
		}
		binary.Write(w, binary.BigEndian, byte((value&0x7F)|0x80))
		value = int(uint32(value) >> 7)
	}
}

func readVarInt(r io.Reader) (int, error) {
	var num int
	var shift uint
	for {
		b := make([]byte, 1)
		_, err := r.Read(b)
		if err != nil {
			return 0, err
		}
		num |= int(b[0]&0x7F) << shift
		shift += 7
		if b[0]&0x80 == 0 {
			break
		}
	}
	return num, nil
}

func writeString(w io.Writer, s string) error {
	if err := writeVarInt(w, len(s)); err != nil {
		return err
	}
	_, err := w.Write([]byte(s))
	return err
}

func sendHandshake(conn net.Conn, addr string, port int) error {
	handshake := new(bytes.Buffer)
	writeVarInt(handshake, 0)    // 包ID
	writeVarInt(handshake, -1)   // 协议版本
	writeString(handshake, addr) // 服务器地址
	binary.Write(handshake, binary.BigEndian, uint16(port))
	writeVarInt(handshake, 1) // 下一个状态 (Status)

	packet := new(bytes.Buffer)
	writeVarInt(packet, handshake.Len())
	packet.Write(handshake.Bytes())
	_, err := conn.Write(packet.Bytes())
	return err
}

func sendStatusRequest(conn net.Conn) error {
	request := new(bytes.Buffer)
	writeVarInt(request, 0) // 请求包ID

	packet := new(bytes.Buffer)
	writeVarInt(packet, request.Len())
	packet.Write(request.Bytes())
	_, err := conn.Write(packet.Bytes())
	return err
}

func readStatusResponse(conn net.Conn) ([]byte, error) {
	length, err := readVarInt(conn)
	if err != nil {
		return nil, err
	}

	data := make([]byte, length)
	_, err = io.ReadFull(conn, data)
	if err != nil {
		return nil, err
	}

	r := bytes.NewReader(data)
	packetID, err := readVarInt(r)
	if err != nil || packetID != 0 {
		return nil, fmt.Errorf("invalid packet ID: %d", packetID)
	}

	jsonStr, err := readString(r)
	if err != nil {
		return nil, err
	}

	return []byte(jsonStr), nil
}

// 从输入流中读取一个字符串
func readString(r io.Reader) (string, error) {
	length, err := readVarInt(r)
	if err != nil {
		return "", err
	}

	if length < 0 {
		return "", fmt.Errorf("invalid string length: %d", length)
	}

	data := make([]byte, length)
	_, err = io.ReadFull(r, data)
	if err != nil {
		return "", err
	}

	return string(data), nil
}

func parseDescription(raw json.RawMessage) string {
	// 尝试解析为text component
	var obj struct {
		Text string `json:"text"`
	}
	if json.Unmarshal(raw, &obj) == nil {
		return obj.Text
	}

	// 直接返回原始字符串
	var str string
	if json.Unmarshal(raw, &str) == nil {
		return str
	}

	return string(raw)
}