package fastmap

import (
	"fmt"
	"strings"
	"sync"

	"slices"

	"github.com/goccy/go-json"
)

type Fastmap[K comparable, V any] struct {
	mu    sync.RWMutex
	idx   map[K]int
	store []fastmapValue[K, V]
}

type fastmapValue[K comparable, V any] struct {
	Key   K
	Value V
}

func New[K comparable, V any]() *Fastmap[K, V] {
	return &Fastmap[K, V]{
		idx:   make(map[K]int),
		store: make([]fastmapValue[K, V], 0),
	}
}

func (m *Fastmap[K, V]) Set(key K, value V) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if _, ok := m.idx[key]; ok {
		m.store[m.idx[key]].Value = value
		return
	}
	m.idx[key] = len(m.store)
	m.store = append(m.store, fastmapValue[K, V]{Key: key, Value: value})
}

func (m *Fastmap[K, V]) Get(key K) (value V, ok bool) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	idx, ok := m.idx[key]
	if !ok {
		return
	}
	return m.store[idx].Value, true
}

func (m *Fastmap[K, V]) Delete(key K) {
	m.mu.Lock()
	defer m.mu.Unlock()
	idx, ok := m.idx[key]
	if !ok {
		return
	}
	delete(m.idx, key)
	m.store[idx] = m.store[len(m.store)-1]
	m.store = m.store[:len(m.store)-1]
}

func (m *Fastmap[K, V]) Has(key K) bool {
	m.mu.RLock()
	defer m.mu.RUnlock()
	_, ok := m.idx[key]
	return ok
}

func (m *Fastmap[K, V]) Len() int {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return len(m.idx)
}

func (m *Fastmap[K, V]) GetPage(pageNum int, pageSize int) *Fastmap[K, V] {
	m.mu.RLock()
	defer m.mu.RUnlock()
	start := pageSize * pageNum
	end := start + pageSize
	if end > len(m.store) {
		end = len(m.store)
	}

	returnVal := New[K, V]()
	for i := start; i < end; i++ {
		returnVal.Set(m.store[i].Key, m.store[i].Value)
	}
	return returnVal
}

func (m *Fastmap[K, V]) Clear() {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.idx = make(map[K]int)
	m.store = make([]fastmapValue[K, V], 0)
}

func (m *Fastmap[K, V]) Iter(fn func(key K, value V) bool) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	for _, v := range m.store {
		if !fn(v.Key, v.Value) {
			break
		}
	}
}

func (m *Fastmap[K, V]) StableSortByKey() {
	m.mu.Lock()
	defer m.mu.Unlock()
	slices.SortStableFunc(m.store, func(a, b fastmapValue[K, V]) int {
		aKey := fmt.Sprint(a.Key)
		bKey := fmt.Sprint(b.Key)
		return strings.Compare(aKey, bKey)
	})

	// Update the index map after sorting
	for i, v := range m.store {
		m.idx[v.Key] = i
	}
}

func (m *Fastmap[K, V]) MarshalText() ([]byte, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	var builder strings.Builder
	for _, v := range m.store {
		builder.WriteString(fmt.Sprintf("%v:%v\n", v.Key, v.Value))
	}
	return []byte(builder.String()), nil
}

func (m *Fastmap[K, V]) UnmarshalText(text []byte) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.Clear()
	lines := strings.Split(string(text), "\n")
	for _, line := range lines {
		if line == "" {
			continue
		}
		parts := strings.SplitN(line, ":", 2)
		if len(parts) != 2 {
			return fmt.Errorf("invalid format: %s", line)
		}
		var key K
		var value V
		if _, err := fmt.Sscan(parts[0], &key); err != nil {
			return fmt.Errorf("error parsing key: %v", err)
		}
		if _, err := fmt.Sscan(parts[1], &value); err != nil {
			return fmt.Errorf("error parsing value: %v", err)
		}
		m.Set(key, value)
	}
	return nil
}

func (m *Fastmap[K, V]) MarshalJSON() ([]byte, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	temp := make(map[K]V)
	for _, v := range m.store {
		temp[v.Key] = v.Value
	}
	return json.Marshal(temp)
}

func (m *Fastmap[K, V]) UnmarshalJSON(data []byte) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	temp := make(map[K]V)
	if err := json.Unmarshal(data, &temp); err != nil {
		return err
	}
	m.Clear()
	for k, v := range temp {
		m.Set(k, v)
	}
	return nil
}