// registry.go - Session-scoped tool registry for toolsrv.
package main

import (
	"fmt"
	"sort"
	"sync"

	"ollie/toolsrv"
)

// Registry tracks loaded tools per session.
type Registry struct {
	mu        sync.RWMutex
	sessions  map[string]map[string]toolsrv.ToolInfo
	revisions map[string]uint64
}

// NewRegistry creates a new tool registry.
func NewRegistry() (*Registry, error) {
	r := &Registry{
		sessions:  make(map[string]map[string]toolsrv.ToolInfo),
		revisions: make(map[string]uint64),
	}
	return r, nil
}

// discover scans the tools directory and returns all valid tool infos.
func (r *Registry) discover() map[string]toolsrv.ToolInfo {
	infos := toolsrv.DiscoverTools()
	tools := make(map[string]toolsrv.ToolInfo, len(infos))
	for _, info := range infos {
		tools[info.Name] = info
	}
	return tools
}

// Load loads a tool into a session's registry.
func (r *Registry) Load(sessionID, name string) error {
	tools := r.discover()
	tool, exists := tools[name]
	if !exists {
		return fmt.Errorf("tool not found: %s", name)
	}

	r.mu.Lock()
	defer r.mu.Unlock()

	if r.sessions[sessionID] == nil {
		r.sessions[sessionID] = make(map[string]toolsrv.ToolInfo)
	}

	r.sessions[sessionID][name] = tool
	r.revisions[sessionID]++
	return nil
}

// Unload removes a tool from a session's registry.
func (r *Registry) Unload(sessionID, name string) error {
	r.mu.Lock()
	defer r.mu.Unlock()

	sessionTools, ok := r.sessions[sessionID]
	if !ok {
		return nil
	}

	if _, exists := sessionTools[name]; !exists {
		return nil
	}

	delete(sessionTools, name)
	r.revisions[sessionID]++
	return nil
}

// Loaded returns all tools loaded for a session.
func (r *Registry) Loaded(sessionID string) []toolsrv.ToolInfo {
	r.mu.RLock()
	defer r.mu.RUnlock()

	sessionTools, ok := r.sessions[sessionID]
	if !ok {
		return nil
	}

	var loaded []toolsrv.ToolInfo
	for _, info := range sessionTools {
		loaded = append(loaded, info)
	}
	sort.Slice(loaded, func(i, j int) bool { return loaded[i].Name < loaded[j].Name })
	return loaded
}

// Lookup returns a tool by name from a session's registry.
func (r *Registry) Lookup(sessionID, name string) (toolsrv.ToolInfo, bool) {
	r.mu.RLock()
	defer r.mu.RUnlock()

	sessionTools, ok := r.sessions[sessionID]
	if !ok {
		return toolsrv.ToolInfo{}, false
	}

	tool, exists := sessionTools[name]
	return tool, exists
}

// Revision returns the revision counter for a session's registry.
func (r *Registry) Revision(sessionID string) uint64 {
	r.mu.RLock()
	defer r.mu.RUnlock()
	return r.revisions[sessionID]
}