package grpc

import (
	"sync"
	"time"

	agentv1 "github.com/containarr/server/internal/proto/agentv1"
)

type AgentState struct {
	ID         string
	Hostname   string
	Alias      string
	IPAddress  string
	Arch       string
	OS         string
	LastSeenAt time.Time
	Containers []*agentv1.ContainerInfo

	cmdCh chan *agentv1.ServerMessage
}

type Registry struct {
	mu     sync.RWMutex
	agents map[string]*AgentState
}

func NewRegistry() *Registry {
	return &Registry{agents: make(map[string]*AgentState)}
}

func (r *Registry) Register(id, hostname, alias, ipAddress, arch, os string) *AgentState {
	state := &AgentState{
		ID:        id,
		Hostname:  hostname,
		Alias:     alias,
		IPAddress: ipAddress,
		Arch:      arch,
		OS:        os,
		cmdCh:     make(chan *agentv1.ServerMessage, 16),
	}
	r.mu.Lock()
	r.agents[id] = state
	r.mu.Unlock()
	return state
}

func (r *Registry) Deregister(id string) {
	r.mu.Lock()
	if s, ok := r.agents[id]; ok {
		close(s.cmdCh)
		delete(r.agents, id)
	}
	r.mu.Unlock()
}

func (r *Registry) Get(id string) (*AgentState, bool) {
	r.mu.RLock()
	defer r.mu.RUnlock()
	s, ok := r.agents[id]
	return s, ok
}

func (r *Registry) List() []*AgentState {
	r.mu.RLock()
	defer r.mu.RUnlock()
	out := make([]*AgentState, 0, len(r.agents))
	for _, s := range r.agents {
		out = append(out, s)
	}
	return out
}

func (r *Registry) UpdateContainers(id string, containers []*agentv1.ContainerInfo) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if s, ok := r.agents[id]; ok {
		s.Containers = containers
		s.LastSeenAt = time.Now()
	}
}

// UpdateAlias refreshes the alias for a live agent (called after an admin update).
func (r *Registry) UpdateAlias(id, alias string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if s, ok := r.agents[id]; ok {
		s.Alias = alias
	}
}

func (r *Registry) Send(agentID string, msg *agentv1.ServerMessage) bool {
	r.mu.RLock()
	s, ok := r.agents[agentID]
	r.mu.RUnlock()
	if !ok {
		return false
	}
	select {
	case s.cmdCh <- msg:
		return true
	default:
		return false
	}
}