visionA/local-agent/server/internal/device/manager.go

package device

import (
	"fmt"
	"log"
	"sync"

	"visiona-agent/server/internal/driver"
	"visiona-agent/server/internal/driver/kneron"
	"visiona-agent/server/pkg/logger"
)

// fakeSerialNumber is the placeholder kn_number reported by the Python
// bridge pyusb fallback when the Kneron SDK is unavailable (e.g. macOS
// without the dylib). It is not a real hardware serial, so it must never be
// used as a routing key (multiple devices could collide on it). Treated the
// same as "no serial" (ADR-018 §2.2 / R2).
const fakeSerialNumber = "0x00000000"

type Manager struct {
	registry *DriverRegistry
	sessions map[string]*DeviceSession
	// serialToLocalID maps a device's hardware serial (Kneron kn_number,
	// e.g. "0x1A2B3C4D") to its local synthetic session key ("kl520-0").
	//
	// Why it exists (ADR-018 serial routing): cloud-side requests arrive with
	// an identifier that never matched the local sessions key (the cloud used
	// its own DB UUID). The serial is the only identifier stable across both
	// layers, so GetDevice does a dual lookup: sessions[id] first (backward
	// compatible with local synthetic IDs), then serialToLocalID[id].
	// Empty and fake ("0x00000000") serials are never indexed.
	serialToLocalID map[string]string
	eventBus        chan DeviceEvent
	scriptPath      string
	logBroadcaster  *logger.Broadcaster
	mu              sync.RWMutex

	// detectFn / newDriverFn are seams so Rescan behavior (unplug / slot
	// swap of positional synthetic IDs) can be tested without spawning the
	// Python bridge. Production wiring (set in NewManager) uses the kneron
	// implementations.
	detectFn    func(scriptPath string) []driver.DeviceInfo
	newDriverFn func(info driver.DeviceInfo, scriptPath string) driver.DeviceDriver
}

func NewManager(registry *DriverRegistry, scriptPath string) *Manager {
	return &Manager{
		registry:        registry,
		sessions:        make(map[string]*DeviceSession),
		serialToLocalID: make(map[string]string),
		eventBus:        make(chan DeviceEvent, 100),
		scriptPath:      scriptPath,
		detectFn:        kneron.DetectDevices,
		newDriverFn: func(info driver.DeviceInfo, scriptPath string) driver.DeviceDriver {
			return kneron.NewKneronDriver(info, scriptPath)
		},
	}
}

// SetLogBroadcaster attaches a log broadcaster so that Kneron driver
// and bridge logs are forwarded to the frontend.
func (m *Manager) SetLogBroadcaster(b *logger.Broadcaster) {
	m.logBroadcaster = b
	// Also set on any already-registered kneron drivers.
	m.mu.RLock()
	defer m.mu.RUnlock()
	for _, s := range m.sessions {
		if kd, ok := s.Driver.(*kneron.KneronDriver); ok {
			kd.SetLogBroadcaster(b)
		}
	}
}

// newSessionDriverLocked builds a driver for a detected device and attaches
// the log broadcaster when the concrete driver supports it. Caller must hold
// m.mu (write lock) — it only reads m.logBroadcaster/scriptPath, but is kept
// locked-only for consistency with its callers.
func (m *Manager) newSessionDriverLocked(info driver.DeviceInfo) driver.DeviceDriver {
	d := m.newDriverFn(info, m.scriptPath)
	if m.logBroadcaster != nil {
		if kd, ok := d.(*kneron.KneronDriver); ok {
			kd.SetLogBroadcaster(m.logBroadcaster)
		}
	}
	return d
}

func (m *Manager) Start() {
	// Detect real Kneron devices (KL520, KL720, etc.) via Python bridge.
	devices := m.detectFn(m.scriptPath)
	if len(devices) == 0 {
		log.Println("No Kneron devices detected")
		return
	}

	m.mu.Lock()
	defer m.mu.Unlock()
	for _, info := range devices {
		m.sessions[info.ID] = NewSession(m.newSessionDriverLocked(info))
		log.Printf("Registered Kneron device: %s (%s, type=%s)", info.Name, info.ID, info.Type)
	}
	m.rebuildSerialIndexLocked()
}

// serialIdentity normalizes a serial for identity comparison: empty and the
// fake pyusb-fallback serial both mean "no identity" (ADR-018 R2), so they
// compare equal to each other and never equal to a real kn_number.
func serialIdentity(serial string) string {
	if serial == fakeSerialNumber {
		return ""
	}
	return serial
}

// Rescan re-detects connected Kneron devices. New devices are registered and
// removed devices are cleaned up. For an existing synthetic ID the session is
// kept only when the detected serial identity matches the session's; on a
// mismatch the session is replaced with a fresh driver built from the
// detected info.
//
// Why the identity check (WP-0 review Minor #1): synthetic IDs are
// positional ("kl520-0" = first KL520 found). With multiple same-chip
// dongles, unplugging or re-ordering shifts which physical device occupies a
// slot. Keeping the old session would make rebuildSerialIndexLocked() index
// a stale serial → the serial route would hit a removed/relocated device.
// The index must always reflect what detection currently sees.
func (m *Manager) Rescan() []driver.DeviceInfo {
	detected := m.detectFn(m.scriptPath)

	// Build a set of detected device IDs.
	detectedIDs := make(map[string]driver.DeviceInfo, len(detected))
	for _, info := range detected {
		detectedIDs[info.ID] = info
	}

	// Drivers to disconnect after releasing the lock: a real KneronDriver
	// Disconnect can take seconds, and holding the write lock would block
	// every GetDevice (inference / flash / camera routing) meanwhile.
	type staleDriver struct {
		id  string
		drv driver.DeviceDriver
	}
	var stale []staleDriver

	m.mu.Lock()

	// Remove devices that are no longer present.
	for id, s := range m.sessions {
		if _, exists := detectedIDs[id]; !exists {
			log.Printf("Device removed: %s", id)
			stale = append(stale, staleDriver{id: id, drv: s.Driver})
			delete(m.sessions, id)
		}
	}

	// Register newly detected devices and replace sessions whose slot is now
	// occupied by a different physical device.
	for _, info := range detected {
		if s, exists := m.sessions[info.ID]; exists {
			old := s.Driver.Info().SerialNumber
			if serialIdentity(old) == serialIdentity(info.SerialNumber) {
				// Same physical device (or no identity on either side, e.g.
				// the single-dongle no-SDK demo) — keep the live session and
				// its connection state untouched.
				continue
			}
			log.Printf("Device %s changed identity (serial %q -> %q); replacing session", info.ID, old, info.SerialNumber)
			stale = append(stale, staleDriver{id: info.ID, drv: s.Driver})
			delete(m.sessions, info.ID)
		}
		m.sessions[info.ID] = NewSession(m.newSessionDriverLocked(info))
		log.Printf("Registered Kneron device: %s (%s, type=%s)", info.Name, info.ID, info.Type)
	}

	m.rebuildSerialIndexLocked()

	// Snapshot current list while still holding the lock.
	devices := make([]driver.DeviceInfo, 0, len(m.sessions))
	for _, s := range m.sessions {
		devices = append(devices, s.Driver.Info())
	}
	m.mu.Unlock()

	for _, sd := range stale {
		if err := sd.drv.Disconnect(); err != nil {
			log.Printf("WARNING: disconnect of stale device %s failed: %v", sd.id, err)
		}
	}

	return devices
}

// rebuildSerialIndexLocked rebuilds serialToLocalID from the current
// sessions. Caller must hold m.mu (write lock).
//
// Rebuilding (instead of incremental add/delete) keeps the index trivially
// consistent with sessions across Start/Rescan, including device removal.
// Empty serials are skipped; the fake serial "0x00000000" is skipped because
// multiple SDK-less devices report the same value (routing would be
// ambiguous). Real kn_numbers are unique per physical dongle; should a
// duplicate ever appear it is logged and only one entry wins.
func (m *Manager) rebuildSerialIndexLocked() {
	idx := make(map[string]string, len(m.sessions))
	for id, s := range m.sessions {
		serial := s.Driver.Info().SerialNumber
		if serial == "" || serial == fakeSerialNumber {
			continue
		}
		if prev, dup := idx[serial]; dup {
			log.Printf("WARNING: duplicate device serial %s (devices %s and %s); serial routing keeps one entry only", serial, prev, id)
			continue
		}
		idx[serial] = id
	}
	m.serialToLocalID = idx
}

func (m *Manager) ListDevices() []driver.DeviceInfo {
	m.mu.RLock()
	defer m.mu.RUnlock()
	devices := make([]driver.DeviceInfo, 0, len(m.sessions))
	for _, s := range m.sessions {
		devices = append(devices, s.Driver.Info())
	}
	return devices
}

// GetDevice resolves a device session by identifier with a dual lookup
// (ADR-018 serial routing, additive and backward compatible):
//
//  1. sessions[id] — the local synthetic key ("kl520-0"); every existing
//     caller keeps working unchanged.
//  2. serialToLocalID[id] — the hardware serial (kn_number, "0x1A2B3C4D");
//     lets cloud-proxied requests (flash / inference / camera / connect /
//     disconnect all converge here) address a physical dongle by serial.
func (m *Manager) GetDevice(id string) (*DeviceSession, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	if s, ok := m.sessions[id]; ok {
		return s, nil
	}
	if localID, ok := m.serialToLocalID[id]; ok {
		if s, ok := m.sessions[localID]; ok {
			return s, nil
		}
	}
	return nil, fmt.Errorf("device not found: %s", id)
}

func (m *Manager) Connect(id string) error {
	s, err := m.GetDevice(id)
	if err != nil {
		return err
	}
	if err := s.Driver.Connect(); err != nil {
		return err
	}
	m.eventBus <- DeviceEvent{Event: "updated", Device: s.Driver.Info()}
	return nil
}

func (m *Manager) Disconnect(id string) error {
	s, err := m.GetDevice(id)
	if err != nil {
		return err
	}
	if err := s.Driver.Disconnect(); err != nil {
		return err
	}
	m.eventBus <- DeviceEvent{Event: "updated", Device: s.Driver.Info()}
	return nil
}

func (m *Manager) Events() <-chan DeviceEvent {
	return m.eventBus
}