package bedrock

import (
	"bytes"
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"math/rand"
	"net"
	"strconv"
	"strings"
	"time"

	"github.com/VexoraDevelopment/gamestatus/info"
)

type Options struct {
	Timeout time.Duration

	EnableSRV bool

	Resolver *net.Resolver
}

type Client struct {
	opt Options
}

func New(opt Options) *Client {
	if opt.Timeout <= 0 {
		opt.Timeout = 2 * time.Second
	}
	if opt.Resolver == nil {
		opt.Resolver = net.DefaultResolver
	}
	return &Client{opt: opt}
}

func (c *Client) Connect(address string, port int) (*info.ServerInfo, error) {
	if port == 0 {
		port = 19132
	}

	host := address
	p := port

	if c.opt.EnableSRV && net.ParseIP(host) == nil {
		if rh, rp, err := c.resolveMinecraftSRV(host); err == nil && rh != "" {
			host = rh
			if rp != 0 {
				p = rp
			}
		}
	}

	if info, err := c.QueryFull(host, p); err == nil {
		return info, nil
	}

	return c.PingRaknet(host, p)
}

func (c *Client) QueryFull(address string, port int) (*info.ServerInfo, error) {
	conn, _, err := c.dialUDP(address, port)
	if err != nil {
		return nil, err
	}
	defer conn.Close()

	chalPayload, err := c.writeQueryPacket(conn, 0x09, nil)
	if err != nil {
		return nil, err
	}

	chalStr := strings.TrimRight(string(chalPayload), "\x00\r\n\t ")
	chalInt, err := strconv.Atoi(chalStr)
	if err != nil {
		return nil, fmt.Errorf("query challenge parse failed: %w (payload=%q)", err, chalStr)
	}

	chal := make([]byte, 4)
	binary.BigEndian.PutUint32(chal, uint32(chalInt))

	appendData := make([]byte, 0, 8)
	appendData = append(appendData, chal...)
	appendData = append(appendData, 0x00, 0x00, 0x00, 0x00)

	data, err := c.writeQueryPacket(conn, 0x00, appendData)
	if err != nil {
		return nil, err
	}

	if len(data) < 11 {
		return nil, errors.New("query full stat response too short")
	}
	payload := data[11:]

	marker := []byte{0x00, 0x00, 0x01, 'p', 'l', 'a', 'y', 'e', 'r', '_', 0x00, 0x00}
	parts := bytes.SplitN(payload, marker, 2)
	if len(parts) != 2 {
		return nil, errors.New("query full stat parse failed: marker not found")
	}

	kvPart := parts[0]
	playersPart := parts[1]

	if len(playersPart) >= 2 {
		playersPart = playersPart[:len(playersPart)-2]
	}

	raw := map[string]string{}
	kvs := bytes.Split(kvPart, []byte{0x00})

	for i := 0; i+1 < len(kvs); i += 2 {
		k := string(kvs[i])
		v := string(kvs[i+1])
		if k == "" {
			continue
		}
		raw[k] = v
	}

	var players []string
	if len(playersPart) > 0 {
		pp := bytes.Split(playersPart, []byte{0x00})
		for _, b := range pp {
			s := string(b)
			if s != "" {
				players = append(players, s)
			}
		}
	}

	ip := resolveIP(address)
	srvInfo := &info.ServerInfo{
		Edition:  info.EditionBedrock,
		Address:  address,
		IP:       ip,
		Port:     intFromMap(raw, "hostport", port),
		MOTD:     raw["hostname"],
		Version:  raw["version"],
		Game:     raw["game_id"],
		Map:      raw["map"],
		GameMode: raw["gametype"],
		Online:   intFromMap(raw, "numplayers", 0),
		Slots:    intFromMap(raw, "maxplayers", 0),
		Players:  players,
		Query:    true,
		RawQuery: raw,
	}

	if pl := raw["plugins"]; pl != "" {
		parts := strings.SplitN(pl, ": ", 2)
		srvInfo.Software = parts[0]
		if len(parts) == 2 {
			p2 := strings.Split(parts[1], "; ")
			if len(p2) > 1 {
				srvInfo.Plugins = p2
			} else if len(p2) == 1 && p2[0] != "" {
				srvInfo.Plugins = []string{p2[0]}
			}
		}
	} else if eng := raw["server_engine"]; eng != "" {
		srvInfo.Software = eng
	}

	return srvInfo, nil
}

func (c *Client) PingRaknet(address string, port int) (*info.ServerInfo, error) {
	conn, _, err := c.dialUDP(address, port)
	if err != nil {
		return nil, err
	}
	defer conn.Close()

	req := buildUnconnectedPing()

	if err := conn.SetDeadline(time.Now().Add(c.opt.Timeout)); err != nil {
		return nil, err
	}
	if _, err := conn.Write(req); err != nil {
		return nil, err
	}

	buf := make([]byte, 4096)
	n, err := conn.Read(buf)
	if err != nil {
		return nil, err
	}
	resp := buf[:n]

	if len(resp) == 0 || resp[0] != 0x1C {
		return nil, errors.New("raknet pong not received (unexpected response)")
	}
	if len(resp) < 35 {
		return nil, errors.New("raknet pong too short")
	}

	s := string(resp[35:])
	s = strings.TrimRight(s, "\x00\r\n\t ")

	parts := strings.Split(s, ";")
	motd := safeIndex(parts, 1)
	version := safeIndex(parts, 3)
	online := atoiSafe(safeIndex(parts, 4))
	slots := atoiSafe(safeIndex(parts, 5))

	ip := resolveIP(address)
	srvInfo := &info.ServerInfo{
		Edition: info.EditionBedrock,
		Address: address,
		IP:      ip,
		Port:    port,
		MOTD:    motd,
		Version: version,
		Online:  online,
		Slots:   slots,
		Query:   false,
	}
	return srvInfo, nil
}

func (c *Client) dialUDP(host string, port int) (*net.UDPConn, *net.UDPAddr, error) {
	raddr, err := net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", host, port))
	if err != nil {
		return nil, nil, err
	}
	conn, err := net.DialUDP("udp", nil, raddr)
	if err != nil {
		return nil, nil, err
	}
	if err := conn.SetDeadline(time.Now().Add(c.opt.Timeout)); err != nil {
		_ = conn.Close()
		return nil, nil, err
	}
	return conn, raddr, nil
}

func (c *Client) writeQueryPacket(conn *net.UDPConn, command byte, appendData []byte) ([]byte, error) {
	header := []byte{0xFE, 0xFD, command, 0x01, 0x02, 0x03, 0x04}
	packet := make([]byte, 0, len(header)+len(appendData))
	packet = append(packet, header...)
	packet = append(packet, appendData...)

	if err := conn.SetDeadline(time.Now().Add(c.opt.Timeout)); err != nil {
		return nil, err
	}
	if _, err := conn.Write(packet); err != nil {
		return nil, err
	}

	buf := make([]byte, 4096)
	n, err := conn.Read(buf)
	if err != nil {
		return nil, err
	}
	resp := buf[:n]
	if len(resp) < 5 || resp[0] != command {
		return nil, errors.New("query response validation failed")
	}
	return resp[5:], nil
}

func buildUnconnectedPing() []byte {
	req := make([]byte, 0, 1+8+16+8)
	req = append(req, 0x01)

	pingID := rand.New(rand.NewSource(time.Now().UnixNano())).Uint64()
	tmp := make([]byte, 8)
	binary.BigEndian.PutUint64(tmp, pingID)
	req = append(req, tmp...)

	magic := []byte{0x00, 0xFF, 0xFF, 0x00, 0xFE, 0xFE, 0xFE, 0xFE, 0xFD, 0xFD, 0xFD, 0xFD, 0x12, 0x34, 0x56, 0x78}
	req = append(req, magic...)

	binary.BigEndian.PutUint64(tmp, 0)
	req = append(req, tmp...)
	return req
}

func (c *Client) resolveMinecraftSRV(domain string) (string, int, error) {
	ctx, cancel := context.WithTimeout(context.Background(), c.opt.Timeout)
	defer cancel()

	type srvQ struct{ service, proto string }
	for _, q := range []srvQ{
		{service: "minecraft", proto: "udp"},
		{service: "minecraft", proto: "tcp"},
	} {
		_, addrs, err := c.opt.Resolver.LookupSRV(ctx, q.service, q.proto, domain)
		if err != nil || len(addrs) == 0 {
			continue
		}
		target := strings.TrimSuffix(addrs[0].Target, ".")
		port := int(addrs[0].Port)
		if target != "" {
			return target, port, nil
		}
	}
	return "", 0, errors.New("no minecraft SRV record found")
}

func resolveIP(host string) string {
	if ip := net.ParseIP(host); ip != nil {
		return ip.String()
	}
	ips, err := net.LookupIP(host)
	if err != nil || len(ips) == 0 {
		return ""
	}
	return ips[0].String()
}

func safeIndex(a []string, i int) string {
	if i < 0 || i >= len(a) {
		return ""
	}
	return a[i]
}

func atoiSafe(s string) int {
	n, _ := strconv.Atoi(strings.TrimSpace(s))
	return n
}

func intFromMap(m map[string]string, key string, def int) int {
	if m == nil {
		return def
	}
	v := strings.TrimSpace(m[key])
	if v == "" {
		return def
	}
	n, err := strconv.Atoi(v)
	if err != nil {
		return def
	}
	return n
}