byre/internal/store/lead.go

package store

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"time"

	"g.fc.run/ai/byre/internal/api"
)

type LeadInputKind string

const (
	LeadMessage LeadInputKind = "message"
	LeadSignal  LeadInputKind = "signal"
)

type LeadInput struct {
	ID        int64
	Owner     string
	Kind      LeadInputKind
	Actor     string
	Text      string
	Images    []string
	SendID    string
	Client    string
	CreatedAt time.Time
}

func (s *Store) AddLeadInput(ctx context.Context, in LeadInput) error {
	images, err := json.Marshal(append([]string{}, in.Images...))
	if err != nil {
		return err
	}
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		_, err := tx.ExecContext(ctx, `INSERT INTO lead_inputs (owner, kind, actor, text, images, send_id, client_id, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?)`,
			in.Owner, string(in.Kind), in.Actor, in.Text, string(images), in.SendID, in.Client, now())
		return err
	})
}

func (s *Store) LeadInputs(ctx context.Context, owner string) ([]LeadInput, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT id, owner, kind, actor, text, images, send_id, client_id, created_at FROM lead_inputs WHERE owner = ? ORDER BY id`, owner)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	inputs := []LeadInput{}
	for rows.Next() {
		in := LeadInput{}
		kind := ""
		images := ""
		createdAt := ""
		err = rows.Scan(&in.ID, &in.Owner, &kind, &in.Actor, &in.Text, &images, &in.SendID, &in.Client, &createdAt)
		if err != nil {
			return nil, err
		}
		in.Kind = LeadInputKind(kind)
		err = json.Unmarshal([]byte(images), &in.Images)
		if err != nil {
			return nil, err
		}
		in.CreatedAt, err = parseTime(createdAt)
		if err != nil {
			return nil, err
		}
		inputs = append(inputs, in)
	}
	return inputs, rows.Err()
}

func (s *Store) RestoreLeadInputs(ctx context.Context, inputs []LeadInput) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		for _, in := range inputs {
			images, err := json.Marshal(append([]string{}, in.Images...))
			if err != nil {
				return err
			}
			_, err = tx.ExecContext(ctx, `INSERT INTO lead_inputs (id, owner, kind, actor, text, images, send_id, client_id, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`,
				in.ID, in.Owner, string(in.Kind), in.Actor, in.Text, string(images), in.SendID, in.Client, formatTime(in.CreatedAt))
			if err != nil {
				return err
			}
		}
		return nil
	})
}

type LeadDelivery struct {
	Owner string
	First int64
	Last  int64
}

func leadDeliveries(ctx context.Context, tx *sql.Tx) ([]LeadDelivery, error) {
	rows, err := tx.QueryContext(ctx, `
		SELECT i.owner, min(i.id), max(i.id) FROM lead_inputs i
		WHERE (SELECT l.state FROM agent_sessions l WHERE l.owner = i.owner AND l.state != 'completed' ORDER BY l.created_at DESC LIMIT 1) = 'idle'
		GROUP BY i.owner`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	found := []LeadDelivery{}
	for rows.Next() {
		d := LeadDelivery{}
		err = rows.Scan(&d.Owner, &d.First, &d.Last)
		if err != nil {
			return nil, err
		}
		found = append(found, d)
	}
	return found, rows.Err()
}

func (s *Store) DeleteLeadInputs(ctx context.Context, ids []int64) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		for _, id := range ids {
			_, err := tx.ExecContext(ctx, `DELETE FROM lead_inputs WHERE id = ?`, id)
			if err != nil {
				return err
			}
		}
		return nil
	})
}

func (s *Store) AddLeadWake(ctx context.Context, owner, taskID string) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		active := false
		err := tx.QueryRowContext(ctx, `SELECT status = ? FROM tasks WHERE id = ?`, string(api.TaskStatusActive), taskID).Scan(&active)
		if errors.Is(err, sql.ErrNoRows) || (err == nil && !active) {
			return ErrNotFound
		}
		if err != nil {
			return err
		}
		_, err = tx.ExecContext(ctx, `INSERT OR IGNORE INTO lead_wakes (owner, task_id, created_at) VALUES (?, ?, ?)`, owner, taskID, now())
		return err
	})
}

func (s *Store) TakeLeadWakes(ctx context.Context, taskID string) ([]string, error) {
	owners := []string{}
	err := s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		rows, err := tx.QueryContext(ctx, `DELETE FROM lead_wakes WHERE task_id = ? RETURNING owner`, taskID)
		if err != nil {
			return err
		}
		defer rows.Close()
		for rows.Next() {
			owner := ""
			err = rows.Scan(&owner)
			if err != nil {
				return err
			}
			owners = append(owners, owner)
		}
		return rows.Err()
	})
	return owners, err
}

func (s *Store) GetFocus(ctx context.Context, login string) (api.Focus, error) {
	f := api.Focus{}
	raw := ""
	err := s.db.QueryRowContext(ctx, `SELECT focus FROM users WHERE login = ?`, login).Scan(&raw)
	if errors.Is(err, sql.ErrNoRows) {
		return f, ErrNotFound
	}
	if err != nil {
		return f, err
	}
	err = json.Unmarshal([]byte(raw), &f)
	return f, err
}

func (s *Store) SetFocus(ctx context.Context, login string, f api.Focus, actor string) error {
	raw, err := json.Marshal(f)
	if err != nil {
		return err
	}
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		_, err := tx.ExecContext(ctx, `UPDATE users SET focus = ? WHERE login = ?`, string(raw), login)
		if err != nil {
			return err
		}
		return emit(api.Event{Kind: "focus.updated", Owner: &login, Actor: actor, Payload: ToPayload(f)})
	})
}