internal/memory/rust_client.go

// Package memory provides a client for the Rust memory service.
package memory

import (
	"bufio"
	"context"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"sync"
	"time"

	lz4 "github.com/pierrec/lz4/v4"
)

const (
	// Protocol constants
	magicHeader     uint32 = 0x4D454D41 // "MEMA"
	protocolVersion uint8  = 1

	// Message types
	msgStore            uint8 = 0x01
	msgRetrieve         uint8 = 0x02
	msgUpdate           uint8 = 0x03
	msgDelete           uint8 = 0x04
	msgBatchStore       uint8 = 0x05
	msgGetContextWindow uint8 = 0x10
	msgConsolidate      uint8 = 0x11
	msgSubscribe        uint8 = 0x20
	msgUnsubscribe      uint8 = 0x21
	msgMemoryEvent      uint8 = 0x22
	msgApplyTTL         uint8 = 0x30
	msgDetectConflicts  uint8 = 0x31
	msgResolveConflict  uint8 = 0x32
	msgSuccess          uint8 = 0x80
	msgError            uint8 = 0x81
	msgPartial          uint8 = 0x82

	// Flags
	flagCompressed uint8 = 0x01
	flagChecksum   uint8 = 0x02
	flagEncrypted  uint8 = 0x04
)

// FrameHeader represents the binary protocol frame header
type FrameHeader struct {
	Magic      uint32
	Version    uint8
	Type       uint8
	Flags      uint8
	MessageID  [16]byte
	BodyLength uint64
}

// RustMemoryClient connects to the Rust memory service
type RustMemoryClient struct {
	mu sync.Mutex

	// Connection configuration
	addr           string
	port           int
	compression    bool
	requestTimeout time.Duration

	// Connection pool
	pool *connectionPool

	// Fallback to local backend if Rust service unavailable
	fallback    MemoryManager
	useFallback bool
}

// RustClientConfig configures the Rust client
type RustClientConfig struct {
	Host           string
	Port           int
	Compression    bool
	PoolSize       int
	RequestTimeout time.Duration
	Fallback       MemoryManager
}

// NewRustMemoryClient creates a new Rust memory client
func NewRustMemoryClient(config RustClientConfig) *RustMemoryClient {
	if config.PoolSize == 0 {
		config.PoolSize = 10
	}
	if config.RequestTimeout == 0 {
		config.RequestTimeout = 5 * time.Second
	}

	client := &RustMemoryClient{
		addr:           config.Host,
		port:           config.Port,
		compression:    config.Compression,
		requestTimeout: config.RequestTimeout,
		fallback:       config.Fallback,
		useFallback:    false,
	}

	// Initialize connection pool
	client.pool = newConnectionPool(config.PoolSize, func() (net.Conn, error) {
		return net.DialTimeout("tcp",
			fmt.Sprintf("%s:%d", config.Host, config.Port),
			config.RequestTimeout)
	})

	// Check if Rust service is available
	if err := client.healthCheck(); err != nil {
		client.useFallback = true
	}

	return client
}

// healthCheck verifies connection to Rust service
func (c *RustMemoryClient) healthCheck() error {
	conn, err := c.pool.get()
	if err != nil {
		return err
	}
	defer c.pool.put(conn)

	// Simple ping/pong would go here
	return nil
}

// Store persists a new memory entry
func (c *RustMemoryClient) Store(ctx context.Context, entry *MemoryEntry) error {
	if c.useFallback && c.fallback != nil {
		return c.fallback.Store(ctx, entry)
	}

	conn, err := c.pool.get()
	if err != nil {
		if c.fallback != nil {
			return c.fallback.Store(ctx, entry)
		}
		return err
	}
	defer c.pool.put(conn)

	// Serialize entry
	body, err := serializeMemoryEntry(entry)
	if err != nil {
		return err
	}

	// Build and send frame
	header := c.buildHeader(msgStore, body)
	if err := c.writeFrame(conn, header, body); err != nil {
		return err
	}

	// Read response
	respHeader, respBody, err := c.readFrame(conn)
	if err != nil {
		return err
	}

	if respHeader.Type == msgError {
		return fmt.Errorf("rust service error: %s", string(respBody))
	}

	return nil
}

// BatchStore stores multiple entries efficiently
func (c *RustMemoryClient) BatchStore(ctx context.Context, entries []*MemoryEntry) error {
	if c.useFallback && c.fallback != nil {
		return c.fallback.BatchStore(ctx, entries)
	}

	conn, err := c.pool.get()
	if err != nil {
		if c.fallback != nil {
			return c.fallback.BatchStore(ctx, entries)
		}
		return err
	}
	defer c.pool.put(conn)

	// Serialize entries batch
	body, err := serializeMemoryEntries(entries)
	if err != nil {
		return err
	}

	header := c.buildHeader(msgBatchStore, body)
	if err := c.writeFrame(conn, header, body); err != nil {
		return err
	}

	respHeader, respBody, err := c.readFrame(conn)
	if err != nil {
		return err
	}

	if respHeader.Type == msgError {
		return fmt.Errorf("rust service error: %s", string(respBody))
	}

	return nil
}

// Retrieve searches for relevant memories
func (c *RustMemoryClient) Retrieve(ctx context.Context, query *MemoryQuery) ([]*MemoryEntry, error) {
	if c.useFallback && c.fallback != nil {
		return c.fallback.Retrieve(ctx, query)
	}

	conn, err := c.pool.get()
	if err != nil {
		if c.fallback != nil {
			return c.fallback.Retrieve(ctx, query)
		}
		return nil, err
	}
	defer c.pool.put(conn)

	// Serialize query
	body, err := serializeMemoryQuery(query)
	if err != nil {
		return nil, err
	}

	header := c.buildHeader(msgRetrieve, body)
	if err := c.writeFrame(conn, header, body); err != nil {
		return nil, err
	}

	respHeader, respBody, err := c.readFrame(conn)
	if err != nil {
		return nil, err
	}

	if respHeader.Type == msgError {
		return nil, fmt.Errorf("rust service error: %s", string(respBody))
	}

	return deserializeMemoryEntries(respBody)
}

// Update modifies an existing memory entry
func (c *RustMemoryClient) Update(ctx context.Context, id string, updates map[string]interface{}) error {
	if c.useFallback && c.fallback != nil {
		return c.fallback.Update(ctx, id, updates)
	}

	// TODO: Implement wire protocol for updates
	return fmt.Errorf("update not implemented for Rust client")
}

// Delete removes a memory entry
func (c *RustMemoryClient) Delete(ctx context.Context, id string) error {
	if c.useFallback && c.fallback != nil {
		return c.fallback.Delete(ctx, id)
	}

	conn, err := c.pool.get()
	if err != nil {
		if c.fallback != nil {
			return c.fallback.Delete(ctx, id)
		}
		return err
	}
	defer c.pool.put(conn)

	body := []byte(id)
	header := c.buildHeader(msgDelete, body)
	if err := c.writeFrame(conn, header, body); err != nil {
		return err
	}

	respHeader, respBody, err := c.readFrame(conn)
	if err != nil {
		return err
	}

	if respHeader.Type == msgError {
		return fmt.Errorf("rust service error: %s", string(respBody))
	}

	return nil
}

// DeleteUserData removes all data for a user (GDPR compliance)
func (c *RustMemoryClient) DeleteUserData(ctx context.Context, userID string) error {
	if c.useFallback && c.fallback != nil {
		return c.fallback.DeleteUserData(ctx, userID)
	}

	// TODO: Implement wire protocol for user data deletion
	return fmt.Errorf("delete user data not implemented for Rust client")
}

// GetContextWindow retrieves recent conversation context
func (c *RustMemoryClient) GetContextWindow(ctx context.Context, sessionID string, maxTurns int) ([]*MemoryEntry, error) {
	if c.useFallback && c.fallback != nil {
		return c.fallback.GetContextWindow(ctx, sessionID, maxTurns)
	}

	conn, err := c.pool.get()
	if err != nil {
		if c.fallback != nil {
			return c.fallback.GetContextWindow(ctx, sessionID, maxTurns)
		}
		return nil, err
	}
	defer c.pool.put(conn)

	// Serialize request
	body := make([]byte, len(sessionID)+4)
	copy(body, sessionID)
	binary.BigEndian.PutUint32(body[len(sessionID):], uint32(maxTurns))

	header := c.buildHeader(msgGetContextWindow, body)
	if err := c.writeFrame(conn, header, body); err != nil {
		return nil, err
	}

	respHeader, respBody, err := c.readFrame(conn)
	if err != nil {
		return nil, err
	}

	if respHeader.Type == msgError {
		return nil, fmt.Errorf("rust service error: %s", string(respBody))
	}

	return deserializeMemoryEntries(respBody)
}

// ConsolidateMemory migrates short-term to long-term memory
func (c *RustMemoryClient) ConsolidateMemory(ctx context.Context, sessionID string) error {
	if c.useFallback && c.fallback != nil {
		return c.fallback.ConsolidateMemory(ctx, sessionID)
	}

	conn, err := c.pool.get()
	if err != nil {
		if c.fallback != nil {
			return c.fallback.ConsolidateMemory(ctx, sessionID)
		}
		return err
	}
	defer c.pool.put(conn)

	body := []byte(sessionID)
	header := c.buildHeader(msgConsolidate, body)
	if err := c.writeFrame(conn, header, body); err != nil {
		return err
	}

	respHeader, respBody, err := c.readFrame(conn)
	if err != nil {
		return err
	}

	if respHeader.Type == msgError {
		return fmt.Errorf("rust service error: %s", string(respBody))
	}

	return nil
}

// ApplyTTL removes expired memories and returns count
func (c *RustMemoryClient) ApplyTTL(ctx context.Context) (int64, error) {
	if c.useFallback && c.fallback != nil {
		return c.fallback.ApplyTTL(ctx)
	}

	conn, err := c.pool.get()
	if err != nil {
		if c.fallback != nil {
			return c.fallback.ApplyTTL(ctx)
		}
		return 0, err
	}
	defer c.pool.put(conn)

	header := c.buildHeader(msgApplyTTL, nil)
	if err := c.writeFrame(conn, header, nil); err != nil {
		return 0, err
	}

	respHeader, respBody, err := c.readFrame(conn)
	if err != nil {
		return 0, err
	}

	if respHeader.Type == msgError {
		return 0, fmt.Errorf("rust service error: %s", string(respBody))
	}

	if len(respBody) >= 8 {
		return int64(binary.BigEndian.Uint64(respBody)), nil
	}

	return 0, nil
}

// DetectConflicts finds conflicting memories for a user
func (c *RustMemoryClient) DetectConflicts(ctx context.Context, userID string) ([]Conflict, error) {
	if c.useFallback && c.fallback != nil {
		return c.fallback.DetectConflicts(ctx, userID)
	}

	// TODO: Implement wire protocol for conflict detection
	return nil, nil
}

// ResolveConflict resolves a detected conflict
func (c *RustMemoryClient) ResolveConflict(ctx context.Context, conflictID string, resolution ConflictResolution) error {
	if c.useFallback && c.fallback != nil {
		return c.fallback.ResolveConflict(ctx, conflictID, resolution)
	}

	// TODO: Implement wire protocol for conflict resolution
	return nil
}

// Subscribe creates a channel for real-time memory events
func (c *RustMemoryClient) Subscribe(ctx context.Context, userID string) (<-chan MemoryEvent, error) {
	if c.useFallback && c.fallback != nil {
		return c.fallback.Subscribe(ctx, userID)
	}

	// TODO: Implement streaming subscription
	ch := make(chan MemoryEvent)
	return ch, nil
}

// Unsubscribe removes a subscription
func (c *RustMemoryClient) Unsubscribe(ctx context.Context, userID string) error {
	if c.useFallback && c.fallback != nil {
		return c.fallback.Unsubscribe(ctx, userID)
	}

	return nil
}

// Close cleans up resources
func (c *RustMemoryClient) Close() error {
	c.pool.close()
	if c.fallback != nil {
		return c.fallback.Close()
	}
	return nil
}

// buildHeader creates a frame header
func (c *RustMemoryClient) buildHeader(msgType uint8, body []byte) FrameHeader {
	header := FrameHeader{
		Magic:      magicHeader,
		Version:    protocolVersion,
		Type:       msgType,
		Flags:      0,
		BodyLength: uint64(len(body)),
	}

	if c.compression && len(body) > 1024 {
		header.Flags |= flagCompressed
	}

	return header
}

// writeFrame writes a complete frame to the connection
func (c *RustMemoryClient) writeFrame(conn net.Conn, header FrameHeader, body []byte) error {
	// Set deadline
	conn.SetWriteDeadline(time.Now().Add(c.requestTimeout))

	writer := bufio.NewWriter(conn)

	// Write header (31 bytes)
	if err := binary.Write(writer, binary.BigEndian, header.Magic); err != nil {
		return err
	}
	writer.WriteByte(header.Version)
	writer.WriteByte(header.Type)
	writer.WriteByte(header.Flags)
	writer.Write(header.MessageID[:])
	binary.Write(writer, binary.BigEndian, header.BodyLength)

	// Write body
	if len(body) > 0 {
		writer.Write(body)
	}

	return writer.Flush()
}

// readFrame reads a complete frame from the connection
func (c *RustMemoryClient) readFrame(conn net.Conn) (FrameHeader, []byte, error) {
	// Set deadline
	conn.SetReadDeadline(time.Now().Add(c.requestTimeout))

	reader := bufio.NewReader(conn)
	var header FrameHeader

	// Read header
	if err := binary.Read(reader, binary.BigEndian, &header.Magic); err != nil {
		return header, nil, err
	}
	if header.Magic != magicHeader {
		return header, nil, fmt.Errorf("invalid magic header: %x", header.Magic)
	}

	var err error
	header.Version, err = reader.ReadByte()
	if err != nil {
		return header, nil, err
	}
	header.Type, err = reader.ReadByte()
	if err != nil {
		return header, nil, err
	}
	header.Flags, err = reader.ReadByte()
	if err != nil {
		return header, nil, err
	}
	if _, err := io.ReadFull(reader, header.MessageID[:]); err != nil {
		return header, nil, err
	}
	if err := binary.Read(reader, binary.BigEndian, &header.BodyLength); err != nil {
		return header, nil, err
	}

	// Read body
	body := make([]byte, header.BodyLength)
	if header.BodyLength > 0 {
		if _, err := io.ReadFull(reader, body); err != nil {
			return header, nil, err
		}
	}

	// Decompress if needed
	if header.Flags&flagCompressed != 0 {
		body, err = decompressLZ4(body)
		if err != nil {
			return header, nil, err
		}
	}

	return header, body, nil
}

// Connection pool

type connectionPool struct {
	mu       sync.Mutex
	conns    []net.Conn
	maxConns int
	factory  func() (net.Conn, error)
}

func newConnectionPool(maxConns int, factory func() (net.Conn, error)) *connectionPool {
	return &connectionPool{
		conns:    make([]net.Conn, 0, maxConns),
		maxConns: maxConns,
		factory:  factory,
	}
}

func (p *connectionPool) get() (net.Conn, error) {
	p.mu.Lock()

	if len(p.conns) > 0 {
		conn := p.conns[len(p.conns)-1]
		p.conns = p.conns[:len(p.conns)-1]
		p.mu.Unlock()
		return conn, nil
	}

	p.mu.Unlock()
	return p.factory()
}

func (p *connectionPool) put(conn net.Conn) {
	p.mu.Lock()
	defer p.mu.Unlock()

	if len(p.conns) < p.maxConns {
		p.conns = append(p.conns, conn)
	} else {
		conn.Close()
	}
}

func (p *connectionPool) close() {
	p.mu.Lock()
	defer p.mu.Unlock()

	for _, conn := range p.conns {
		conn.Close()
	}
	p.conns = nil
}

// Serialization helpers using JSON encoding

// serializeMemoryEntry serializes a single memory entry to JSON
func serializeMemoryEntry(entry *MemoryEntry) ([]byte, error) {
	return json.Marshal(entry)
}

// serializeMemoryEntries serializes multiple entries to JSON array
func serializeMemoryEntries(entries []*MemoryEntry) ([]byte, error) {
	return json.Marshal(entries)
}

// serializeMemoryQuery serializes a query to JSON
func serializeMemoryQuery(query *MemoryQuery) ([]byte, error) {
	return json.Marshal(query)
}

// deserializeMemoryEntries deserializes JSON to memory entries
func deserializeMemoryEntries(data []byte) ([]*MemoryEntry, error) {
	if len(data) == 0 {
		return nil, nil
	}

	var entries []*MemoryEntry
	if err := json.Unmarshal(data, &entries); err != nil {
		return nil, fmt.Errorf("failed to deserialize entries: %w", err)
	}
	return entries, nil
}

// decompressLZ4 decompresses LZ4 data
func decompressLZ4(data []byte) ([]byte, error) {
	if len(data) == 0 {
		return data, nil
	}

	// Simple frame format: first 4 bytes = uncompressed length
	if len(data) < 4 {
		return data, nil // Not compressed or invalid
	}

	// Read uncompressed length (little-endian)
	uncompressedLen := int(data[0]) | int(data[1])<<8 | int(data[2])<<16 | int(data[3])<<24

	if uncompressedLen <= 0 || uncompressedLen > 100*1024*1024 { // Max 100MB
		return data, nil // Invalid length, return as-is
	}

	result := make([]byte, uncompressedLen)
	n, err := lz4.UncompressBlock(data[4:], result)
	if err != nil {
		return nil, fmt.Errorf("lz4 decompress failed: %w", err)
	}
	return result[:n], nil
}

// compressLZ4 compresses data using LZ4 (placeholder)
func compressLZ4(data []byte) ([]byte, error) {
	// For now, return as-is
	// In production, use: github.com/pierrec/lz4/v4
	return data, nil
}