package topics

import (
	"sort"
	"strings"
	"sync"
	"time"
)

type Message struct {
	ID        string    `json:"id"`
	Topic     string    `json:"topic"`
	Payload   string    `json:"payload"`
	QOS       byte      `json:"qos"`
	Retained  bool      `json:"retained"`
	Timestamp time.Time `json:"timestamp"`
	Size      int       `json:"size"`
}

type Tree struct {
	mu   sync.RWMutex
	root *node
}

type node struct {
	name         string
	fullName     string
	children     map[string]*node
	messageCount int
	lastMessage  *Message
}

type NodeSnapshot struct {
	Name         string          `json:"name"`
	FullName     string          `json:"fullName"`
	MessageCount int             `json:"messageCount"`
	LastMessage  *Message        `json:"lastMessage,omitempty"`
	Children     []NodeSnapshot  `json:"children"`
}

func NewTree() *Tree {
	return &Tree{
		root: &node{
			name:     "root",
			fullName: "",
			children: make(map[string]*node),
		},
	}
}

func (t *Tree) AddMessage(msg Message) {
	t.mu.Lock()
	defer t.mu.Unlock()

	parts := strings.Split(msg.Topic, "/")
	current := t.root
	for i, part := range parts {
		child, ok := current.children[part]
		if !ok {
			full := strings.Join(parts[:i+1], "/")
			child = &node{
				name:     part,
				fullName: full,
				children: make(map[string]*node),
			}
			current.children[part] = child
		}
		current = child
		current.messageCount++
		if i == len(parts)-1 {
			copyMsg := msg
			current.lastMessage = &copyMsg
		}
	}
}

func (t *Tree) Snapshot() NodeSnapshot {
	t.mu.RLock()
	defer t.mu.RUnlock()

	return snapshotNode(t.root)
}

func snapshotNode(n *node) NodeSnapshot {
	children := make([]NodeSnapshot, 0, len(n.children))
	keys := make([]string, 0, len(n.children))
	for key := range n.children {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	for _, key := range keys {
		child := n.children[key]
		children = append(children, snapshotNode(child))
	}

	var last *Message
	if n.lastMessage != nil {
		copyMsg := *n.lastMessage
		last = &copyMsg
	}

	return NodeSnapshot{
		Name:         n.name,
		FullName:     n.fullName,
		MessageCount: n.messageCount,
		LastMessage:  last,
		Children:     children,
	}
}