package design_twitter

import (
	"container/heap"
	"maps"
	"slices"
)

type TweetHeap [][]Tweet

func (h TweetHeap) Len() int { return len(h) }
func (h TweetHeap) Less(i, j int) bool {
	timeI, timeJ := -1, -1

	if len(h[i]) > 0 {
		timeI = h[i][len(h[i])-1].Time
	}
	if len(h[j]) > 0 {
		timeJ = h[j][len(h[i])-1].Time
	}

	return timeI > timeJ
}
func (h TweetHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }

func (h *TweetHeap) Push(x any) {
	*h = append(*h, x.([]Tweet))
}

func (h *TweetHeap) Pop() any {
	old := *h
	x := old[len(old)-1]
	*h = old[:len(old)-1]
	return x
}

// ----------------------------------

type Tweet struct {
	ID   int
	Time int
}

type Twitter struct {
	tweets  map[int][]Tweet
	follows map[int]map[int]bool
	counter int
}

func Constructor() Twitter {
	return Twitter{
		tweets:  map[int][]Tweet{},
		follows: map[int]map[int]bool{},
		counter: 0,
	}
}

func (t *Twitter) PostTweet(userId int, tweetId int) {
	t.counter++

	t.tweets[userId] = append(t.tweets[userId], Tweet{
		ID:   tweetId,
		Time: t.counter,
	})
}

func (t *Twitter) GetNewsFeed(userId int) []int {
	users := slices.AppendSeq([]int{userId}, maps.Keys(t.follows[userId]))

	hp := TweetHeap{}
	for _, user := range users {
		hp = append(hp, t.tweets[user])
	}
	heap.Init(&hp)

	feed := []int{}
	for range 10 {
		if len(hp[0]) == 0 {
			break
		}

		top, rest := hp[0][len(hp[0])-1], hp[0][:len(hp[0])-1]
		feed = append(feed, top.ID)
		hp[0] = rest

		heap.Fix(&hp, 0)
	}

	return feed
}

func (t *Twitter) Follow(followerId int, followeeId int) {
	if _, ok := t.follows[followerId]; !ok {
		t.follows[followerId] = map[int]bool{}
	}

	t.follows[followerId][followeeId] = true
}

func (t *Twitter) Unfollow(followerId int, followeeId int) {
	if _, ok := t.follows[followerId]; !ok {
		return
	}

	delete(t.follows[followerId], followeeId)
}