package gateway

import (
	"errors"
	"sync"
)

type Client struct {
	ID       string
	SendFunc func(data []byte) error
}

type Gateway struct {
	mu      sync.RWMutex
	clients map[string]*Client  // clientID -> Client
	uidMap  map[string][]string // uid -> []clientID
}

func NewGateway() *Gateway {
	return &Gateway{
		clients: make(map[string]*Client),
		uidMap:  make(map[string][]string),
	}
}

func (g *Gateway) AddClient(c *Client) {
	g.mu.Lock()
	defer g.mu.Unlock()
	g.clients[c.ID] = c
}

func (g *Gateway) RemoveClient(clientID string) {
	g.mu.Lock()
	defer g.mu.Unlock()
	delete(g.clients, clientID)
	for uid, list := range g.uidMap {
		newList := []string{}
		for _, cid := range list {
			if cid != clientID {
				newList = append(newList, cid)
			}
		}
		g.uidMap[uid] = newList
	}
}

func (g *Gateway) SendToAll(msg []byte) {
	g.mu.RLock()
	defer g.mu.RUnlock()
	for _, c := range g.clients {
		_ = c.SendFunc(msg)
	}
}

func (g *Gateway) SendToClient(clientID string, msg []byte) error {
	g.mu.RLock()
	defer g.mu.RUnlock()
	if c, ok := g.clients[clientID]; ok {
		return c.SendFunc(msg)
	}
	return errors.New("client not found")
}

func (g *Gateway) BindUid(clientID, uid string) error {
	g.mu.Lock()
	defer g.mu.Unlock()
	if _, ok := g.clients[clientID]; !ok {
		return errors.New("client not found")
	}
	g.uidMap[uid] = append(g.uidMap[uid], clientID)
	return nil
}

func (g *Gateway) SendToUid(uid string, msg []byte) {
	g.mu.RLock()
	defer g.mu.RUnlock()
	if list, ok := g.uidMap[uid]; ok {
		for _, cid := range list {
			if c, ok := g.clients[cid]; ok {
				_ = c.SendFunc(msg)
			}
		}
	}
}

func (g *Gateway) ListClients() []string {
	g.mu.RLock()
	defer g.mu.RUnlock()
	ids := make([]string, 0, len(g.clients))
	for id := range g.clients {
		ids = append(ids, id)
	}
	return ids
}

func (g *Gateway) ListUids() map[string][]string {
	g.mu.RLock()
	defer g.mu.RUnlock()
	result := make(map[string][]string, len(g.uidMap))
	for uid, list := range g.uidMap {
		result[uid] = append([]string(nil), list...)
	}
	return result
}