// Package mesh 实现 Meshtastic MQTT 数据包的解码与构建。
// 协议实现参考 meshtastic_mqtt_server 工程（MIT License）：
//   - ServiceEnvelope / MeshPacket / Data 使用 protobuf wire 格式手写解析（google.golang.org/protobuf/encoding/protowire）
//   - 加密包使用频道 PSK + AES-CTR，nonce = little-endian(packetID 8B) + little-endian(fromNum 4B) + 4 字节零
//   - channel hash = xor(channelName) ^ xor(psk)
package mesh

import (
	"crypto/aes"
	"crypto/cipher"
	"encoding/base64"
	"encoding/binary"
	"fmt"
	"strings"
	"unicode/utf8"

	"google.golang.org/protobuf/encoding/protowire"
)

// Portnum 常量（与 meshtastic.proto PortNum 一致）。
const (
	PortNumUnknown     = 0
	PortNumTextMessage = 1
	PortNumPosition    = 3
	PortNumNodeInfo    = 4
	PortNumRouting     = 5
	PortNumTelemetry   = 67
	PortNumMapReport   = 73
)

// NodeNumBroadcast 是广播目标节点号。
const NodeNumBroadcast uint32 = 0xffffffff

// defaultMeshtasticPSK 是 Meshtastic 默认频道密钥（PSK 索引 1）。
var defaultMeshtasticPSK = []byte{
	0xD4, 0xF1, 0xBB, 0x3A,
	0x20, 0x29, 0x07, 0x59,
	0xF0, 0xBC, 0xFF, 0xAB,
	0xCF, 0x4E, 0x69, 0x01,
}

// Packet 是解码后的 MeshPacket 关键字段。
type Packet struct {
	From         uint32
	To           uint32
	Channel      uint32
	ID           uint32
	WantAck      bool
	ViaMQTT      bool
	PKIEncrypted bool
	Decoded      *Data
	Encrypted    []byte
}

// Data 是 MeshPacket.decoded 中的 Data 子包。
type Data struct {
	Portnum uint32
	Payload []byte
}

// TextMessage 是一条解码后的文本消息。
type TextMessage struct {
	From uint32
	To   uint32
	Text string
	Hex  string // 非 UTF-8 时的十六进制表示
}

// NodeInfo 是解码后的节点信息（NODEINFO_APP）。
type NodeInfo struct {
	From      uint32
	ID        string
	LongName  string
	ShortName string
	HWModel   uint64
	Role      uint64
	PublicKey []byte
}

// Position 是解码后的位置信息（POSITION_APP）。
type Position struct {
	From      uint32
	Latitude  *float64
	Longitude *float64
	Altitude  *int32
}

// ExpandPSK 展开 Base64 PSK，兼容 Meshtastic 默认索引 PSK 和短 key 补零规则。
func ExpandPSK(pskBase64 string) ([]byte, error) {
	psk, err := base64.StdEncoding.DecodeString(strings.TrimSpace(pskBase64))
	if err != nil {
		return nil, fmt.Errorf("invalid psk: %w", err)
	}
	if len(psk) == 1 {
		idx := psk[0]
		if idx == 0 {
			return []byte{}, nil
		}
		key := append([]byte(nil), defaultMeshtasticPSK...)
		key[len(key)-1] = byte((int(key[len(key)-1]) + int(idx) - 1) & 0xff)
		return key, nil
	}
	if len(psk) > 0 && len(psk) < 16 {
		return append(psk, make([]byte, 16-len(psk))...), nil
	}
	if len(psk) > 16 && len(psk) < 32 {
		return append(psk, make([]byte, 32-len(psk))...), nil
	}
	if len(psk) != 0 && len(psk) != 16 && len(psk) != 24 && len(psk) != 32 {
		return nil, fmt.Errorf("invalid psk length %d: AES keys must be 16, 24, or 32 bytes", len(psk))
	}
	return psk, nil
}

// Decode 解码一个 MQTT 消息 payload（ServiceEnvelope），
// 返回解码结果（*TextMessage / *NodeInfo / *Position / *GenericPacket）与错误。
// 加密包会用频道 PSK 尝试解密；PKI 加密包（无密钥）返回错误。
func Decode(topic string, raw []byte, key []byte) (any, error) {
	env, err := parseServiceEnvelope(raw)
	if err != nil {
		return nil, fmt.Errorf("service envelope: %w", err)
	}
	if env.Packet == nil {
		return nil, fmt.Errorf("no packet in envelope")
	}
	pkt := env.Packet

	if pkt.Decoded == nil && len(pkt.Encrypted) > 0 {
		decoded, status := tryDecrypt(pkt, env.ChannelID, key)
		if decoded == nil {
			return nil, fmt.Errorf("decrypt failed (%s)", status)
		}
		pkt.Decoded = decoded
	}

	if pkt.Decoded == nil {
		return nil, fmt.Errorf("empty packet")
	}

	switch pkt.Decoded.Portnum {
	case PortNumTextMessage:
		text := string(pkt.Decoded.Payload)
		msg := &TextMessage{From: pkt.From, To: pkt.To, Text: text}
		if !utf8.Valid(pkt.Decoded.Payload) {
			msg.Text = ""
			msg.Hex = fmt.Sprintf("%x", pkt.Decoded.Payload)
		}
		return msg, nil
	case PortNumNodeInfo:
		info, err := parseUser(pkt.Decoded.Payload)
		if err != nil {
			return nil, err
		}
		info.From = pkt.From
		return info, nil
	case PortNumPosition:
		pos, err := parsePosition(pkt.Decoded.Payload)
		if err != nil {
			return nil, err
		}
		pos.From = pkt.From
		return pos, nil
	case PortNumTelemetry:
		return &GenericPacket{Portnum: PortNumTelemetry, From: pkt.From, PayloadLen: len(pkt.Decoded.Payload)}, nil
	default:
		return &GenericPacket{Portnum: pkt.Decoded.Portnum, From: pkt.From, PayloadLen: len(pkt.Decoded.Payload)}, nil
	}
}

// GenericPacket 是未详细解码的其他类型包。
type GenericPacket struct {
	Portnum    uint32
	From       uint32
	PayloadLen int
}

// PortnumName 返回 portnum 的可读名称。
func PortnumName(portnum uint32) string {
	if name, ok := portNumNames[portnum]; ok {
		return name
	}
	return fmt.Sprintf("PORTNUM_%d", portnum)
}

// NodeNumToID 把节点号格式化为 !xxxxxxxx。
func NodeNumToID(nodeNum uint32) string {
	return fmt.Sprintf("!%08x", nodeNum)
}

// ParseNodeID 解析 !xxxxxxxx 为节点号。
func ParseNodeID(nodeID string) (uint32, error) {
	value := strings.TrimSpace(nodeID)
	value = strings.TrimPrefix(value, "!")
	if len(value) != 8 {
		return 0, fmt.Errorf("node id must be !xxxxxxxx")
	}
	var num uint32
	if _, err := fmt.Sscanf(value, "%08x", &num); err != nil {
		return 0, fmt.Errorf("invalid node id: %w", err)
	}
	return num, nil
}

// tryDecrypt 用频道 PSK 解密 encrypted 载荷（AES-CTR），并解析出 Data 子包。
func tryDecrypt(pkt *Packet, channelID string, key []byte) (*Data, string) {
	if len(key) == 0 {
		return nil, "psk disables encryption"
	}
	if pkt.Channel != uint32(channelHash(channelID, key)) {
		return nil, "channel hash mismatch"
	}
	plaintext, err := cryptAESCTR(key, pkt.From, pkt.ID, pkt.Encrypted)
	if err != nil {
		return nil, err.Error()
	}
	decoded, err := parseData(plaintext)
	if err != nil {
		return nil, "decrypted bytes are not Data protobuf"
	}
	if decoded.Portnum == PortNumUnknown {
		return nil, "decrypted protobuf has UNKNOWN_APP portnum"
	}
	return decoded, "success"
}

type serviceEnvelope struct {
	Packet    *Packet
	ChannelID string
	GatewayID string
}

func parseServiceEnvelope(payload []byte) (*serviceEnvelope, error) {
	env := &serviceEnvelope{}
	err := walkFields(payload, func(num protowire.Number, typ protowire.Type, value any) error {
		switch num {
		case 1:
			b, ok := value.([]byte)
			if !ok || typ != protowire.BytesType {
				return nil
			}
			packet, err := parseMeshPacket(b)
			if err != nil {
				return err
			}
			env.Packet = packet
		case 2:
			if b, ok := value.([]byte); ok && typ == protowire.BytesType {
				env.ChannelID = string(b)
			}
		case 3:
			if b, ok := value.([]byte); ok && typ == protowire.BytesType {
				env.GatewayID = string(b)
			}
		}
		return nil
	})
	return env, err
}

func parseMeshPacket(payload []byte) (*Packet, error) {
	pkt := &Packet{}
	err := walkFields(payload, func(num protowire.Number, typ protowire.Type, value any) error {
		switch num {
		case 1:
			if v, ok := value.(uint32); ok && typ == protowire.Fixed32Type {
				pkt.From = v
			}
		case 2:
			if v, ok := value.(uint32); ok && typ == protowire.Fixed32Type {
				pkt.To = v
			}
		case 3:
			if v, ok := value.(uint64); ok && typ == protowire.VarintType {
				pkt.Channel = uint32(v)
			}
		case 4:
			if b, ok := value.([]byte); ok && typ == protowire.BytesType {
				decoded, err := parseData(b)
				if err != nil {
					return err
				}
				pkt.Decoded = decoded
			}
		case 5:
			if b, ok := value.([]byte); ok && typ == protowire.BytesType {
				pkt.Encrypted = append([]byte(nil), b...)
			}
		case 6:
			if v, ok := value.(uint32); ok && typ == protowire.Fixed32Type {
				pkt.ID = v
			}
		case 10:
			if v, ok := value.(uint64); ok && typ == protowire.VarintType {
				pkt.WantAck = v != 0
			}
		case 14:
			if v, ok := value.(uint64); ok && typ == protowire.VarintType {
				pkt.ViaMQTT = v != 0
			}
		case 17:
			if v, ok := value.(uint64); ok && typ == protowire.VarintType {
				pkt.PKIEncrypted = v != 0
			}
		}
		return nil
	})
	return pkt, err
}

func parseData(payload []byte) (*Data, error) {
	data := &Data{}
	err := walkFields(payload, func(num protowire.Number, typ protowire.Type, value any) error {
		switch num {
		case 1:
			if v, ok := value.(uint64); ok && typ == protowire.VarintType {
				data.Portnum = uint32(v)
			}
		case 2:
			if b, ok := value.([]byte); ok && typ == protowire.BytesType {
				data.Payload = append([]byte(nil), b...)
			}
		}
		return nil
	})
	return data, err
}

func parseUser(payload []byte) (*NodeInfo, error) {
	user := &NodeInfo{}
	err := walkFields(payload, func(num protowire.Number, typ protowire.Type, value any) error {
		switch num {
		case 1:
			user.ID = stringBytes(typ, value)
		case 2:
			user.LongName = stringBytes(typ, value)
		case 3:
			user.ShortName = stringBytes(typ, value)
		case 5:
			user.HWModel = varintValue(typ, value)
		case 7:
			user.Role = varintValue(typ, value)
		case 8:
			if b, ok := value.([]byte); ok && typ == protowire.BytesType {
				user.PublicKey = append([]byte(nil), b...)
			}
		}
		return nil
	})
	return user, err
}

func parsePosition(payload []byte) (*Position, error) {
	pos := &Position{}
	err := walkFields(payload, func(num protowire.Number, typ protowire.Type, value any) error {
		switch num {
		case 1:
			if v, ok := value.(uint32); ok && typ == protowire.Fixed32Type {
				lat := float64(int32(v)) * 1e-7
				pos.Latitude = &lat
			}
		case 2:
			if v, ok := value.(uint32); ok && typ == protowire.Fixed32Type {
				lon := float64(int32(v)) * 1e-7
				pos.Longitude = &lon
			}
		case 3:
			if typ == protowire.VarintType {
				alt := int32(varintValue(typ, value))
				pos.Altitude = &alt
			}
		}
		return nil
	})
	return pos, err
}

// walkFields 遍历 protobuf wire 字段。
func walkFields(payload []byte, handle func(protowire.Number, protowire.Type, any) error) error {
	for len(payload) > 0 {
		num, typ, n := protowire.ConsumeTag(payload)
		if n < 0 {
			return protowire.ParseError(n)
		}
		payload = payload[n:]

		var value any
		switch typ {
		case protowire.VarintType:
			v, n := protowire.ConsumeVarint(payload)
			if n < 0 {
				return protowire.ParseError(n)
			}
			value = v
			payload = payload[n:]
		case protowire.Fixed32Type:
			v, n := protowire.ConsumeFixed32(payload)
			if n < 0 {
				return protowire.ParseError(n)
			}
			value = v
			payload = payload[n:]
		case protowire.Fixed64Type:
			v, n := protowire.ConsumeFixed64(payload)
			if n < 0 {
				return protowire.ParseError(n)
			}
			value = v
			payload = payload[n:]
		case protowire.BytesType:
			v, n := protowire.ConsumeBytes(payload)
			if n < 0 {
				return protowire.ParseError(n)
			}
			value = v
			payload = payload[n:]
		default:
			n := protowire.ConsumeFieldValue(num, typ, payload)
			if n < 0 {
				return protowire.ParseError(n)
			}
			payload = payload[n:]
		}

		if err := handle(num, typ, value); err != nil {
			return err
		}
	}
	return nil
}

func stringBytes(typ protowire.Type, value any) string {
	if b, ok := value.([]byte); ok && typ == protowire.BytesType {
		return string(b)
	}
	return ""
}

func varintValue(typ protowire.Type, value any) uint64 {
	if v, ok := value.(uint64); ok && typ == protowire.VarintType {
		return v
	}
	return 0
}

func xorHash(data []byte) byte {
	var result byte
	for _, b := range data {
		result ^= b
	}
	return result
}

func channelHash(channelName string, key []byte) byte {
	return xorHash([]byte(channelName)) ^ xorHash(key)
}

// cryptAESCTR 按 Meshtastic nonce 规则执行 AES-CTR；CTR 加密和解密是同一个 XOR 流操作。
func cryptAESCTR(key []byte, fromNum, packetID uint32, input []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	nonce := make([]byte, aes.BlockSize)
	binary.LittleEndian.PutUint64(nonce[0:8], uint64(packetID))
	binary.LittleEndian.PutUint32(nonce[8:12], fromNum)
	output := make([]byte, len(input))
	cipher.NewCTR(block, nonce).XORKeyStream(output, input)
	return output, nil
}

var portNumNames = map[uint32]string{
	0: "UNKNOWN_APP", 1: "TEXT_MESSAGE_APP", 2: "REMOTE_HARDWARE_APP", 3: "POSITION_APP", 4: "NODEINFO_APP",
	5: "ROUTING_APP", 6: "ADMIN_APP", 7: "TEXT_MESSAGE_COMPRESSED_APP", 8: "WAYPOINT_APP", 9: "AUDIO_APP",
	10: "DETECTION_SENSOR_APP", 11: "ALERT_APP", 12: "KEY_VERIFICATION_APP", 13: "REMOTE_SHELL_APP", 32: "REPLY_APP",
	33: "IP_TUNNEL_APP", 34: "PAXCOUNTER_APP", 35: "STORE_FORWARD_PLUSPLUS_APP", 36: "NODE_STATUS_APP", 64: "SERIAL_APP",
	65: "STORE_FORWARD_APP", 66: "RANGE_TEST_APP", 67: "TELEMETRY_APP", 68: "ZPS_APP", 69: "SIMULATOR_APP",
	70: "TRACEROUTE_APP", 71: "NEIGHBORINFO_APP", 72: "ATAK_PLUGIN", 73: "MAP_REPORT_APP", 74: "POWERSTRESS_APP",
	75: "LORAWAN_BRIDGE", 76: "RETICULUM_TUNNEL_APP", 77: "CAYENNE_APP", 78: "ATAK_PLUGIN_V2", 112: "GROUPALARM_APP",
	256: "PRIVATE_APP", 257: "ATAK_FORWARDER", 511: "MAX",
}