dobotai-platform/internal/store/postgres/store.go

package postgres

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"

	"gitea.zheromiku.top/Weknora/dobotai-platform/internal/core"
)

type Store struct {
	pool                       *pgxpool.Pool
	questionTableExportEnabled bool
	faqExportEnabled           bool
}

func (s *Store) SetExternalTableExports(questionEnabled, faqEnabled bool) {
	if s == nil {
		return
	}
	s.questionTableExportEnabled = questionEnabled
	s.faqExportEnabled = faqEnabled
}

func Open(ctx context.Context, databaseURL string) (*Store, error) {
	if databaseURL == "" {
		return nil, errors.New("DATABASE_URL is required")
	}
	pool, err := pgxpool.New(ctx, databaseURL)
	if err != nil {
		return nil, fmt.Errorf("open postgres: %w", err)
	}
	if err := pool.Ping(ctx); err != nil {
		pool.Close()
		return nil, fmt.Errorf("ping postgres: %w", err)
	}
	return &Store{pool: pool}, nil
}
func (s *Store) Close() {
	if s != nil && s.pool != nil {
		s.pool.Close()
	}
}

func (s *Store) BeginTurn(ctx context.Context, turn core.Turn) (core.Turn, bool, error) {
	media := turn.Media
	if media == nil {
		media = []core.MediaAudit{}
	}
	mediaJSON, err := json.Marshal(media)
	if err != nil {
		return core.Turn{}, false, err
	}
	imageURLs := turn.QuestionImageURLs
	if imageURLs == nil {
		imageURLs = []string{}
	}
	imageURLsJSON, err := json.Marshal(imageURLs)
	if err != nil {
		return core.Turn{}, false, err
	}
	tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{})
	if err != nil {
		return core.Turn{}, false, err
	}
	defer func() { _ = tx.Rollback(ctx) }()

	if err = tx.QueryRow(ctx, `INSERT INTO channel_actors
		(id,tenant_key,channel,external_actor_id,actor_type,display_name)
		VALUES($1,$2,$3,$4,'user',$5)
		ON CONFLICT (tenant_key,channel,external_actor_id) DO UPDATE SET
			display_name=CASE WHEN EXCLUDED.display_name<>'' THEN EXCLUDED.display_name ELSE channel_actors.display_name END,
			updated_at=NOW()
		RETURNING id`, turn.ActorID, turn.TenantKey, turn.Channel, turn.ExternalUserID, turn.SenderName).Scan(&turn.ActorID); err != nil {
		return core.Turn{}, false, err
	}

	externalConversation := externalConversationID(turn)
	if externalConversation == "" {
		return core.Turn{}, false, errors.New("external conversation identity is required")
	}
	conversationTitle := ""
	if turn.ConversationType == core.ConversationDirect {
		conversationTitle = strings.TrimSpace(turn.SenderName)
	} else if turn.ConversationType == core.ConversationGroup {
		conversationTitle = strings.TrimSpace(turn.ConversationTitle)
	}
	if err = tx.QueryRow(ctx, `INSERT INTO channel_conversations
		(id,tenant_key,channel,external_conversation_id,conversation_type,title,status)
		VALUES($1,$2,$3,$4,$5,$6,'active')
		ON CONFLICT (tenant_key,channel,external_conversation_id) DO UPDATE SET
			title=CASE WHEN EXCLUDED.title<>'' THEN EXCLUDED.title ELSE channel_conversations.title END,
			updated_at=NOW()
		RETURNING id`, turn.ConversationID, turn.TenantKey, turn.Channel, externalConversation,
		string(turn.ConversationType), conversationTitle).Scan(&turn.ConversationID); err != nil {
		return core.Turn{}, false, err
	}

	if err = tx.QueryRow(ctx, `INSERT INTO messages
		(id,tenant_key,channel,conversation_id,actor_id,external_message_id,correlation_id,
		 direction,role,message_type,text_content,media_json,status,actor_name_snapshot,trigger_person_snapshot,actor_union_id_snapshot,department_snapshot,region_snapshot,safe_metadata_json)
		VALUES($1,$2,$3,$4,$5,$6,$7,'inbound','user',$8,$9,$10,'received',$11,$12,$13,$14,$15,'{}'::JSONB)
		ON CONFLICT (tenant_key,channel,external_message_id) WHERE external_message_id!=''
		DO UPDATE SET updated_at=messages.updated_at
		RETURNING id`, turn.RequestMessageID, turn.TenantKey, turn.Channel, turn.ConversationID,
		turn.ActorID, turn.ExternalMessageID, turn.ID, turn.MessageType, turn.Question, mediaJSON,
		turn.SenderName, turn.TriggerPersonSnapshot, turn.ActorUnionIDSnapshot, turn.DepartmentSnapshot, turn.RegionSnapshot).Scan(&turn.RequestMessageID); err != nil {
		return core.Turn{}, false, err
	}

	command, err := tx.Exec(ctx, `INSERT INTO turns
		(id,tenant_key,channel,external_message_id,session_key,question,status,media_json,request_message_id,agent_id,question_image_urls_json)
		VALUES($1,$2,$3,$4,$5,$6,'processing',$7,$8,$9,$10) ON CONFLICT DO NOTHING`,
		turn.ID, turn.TenantKey, turn.Channel, turn.ExternalMessageID, turn.SessionKey, turn.Question,
		mediaJSON, turn.RequestMessageID, turn.AgentID, imageURLsJSON)
	if err != nil {
		return core.Turn{}, false, err
	}
	if command.RowsAffected() == 1 {
		if err = tx.Commit(ctx); err != nil {
			return core.Turn{}, false, err
		}
		return turn, true, nil
	}
	stored, err := loadTurnWith(ctx, tx, turn.TenantKey, turn.Channel, turn.ExternalMessageID)
	if err != nil {
		return core.Turn{}, false, err
	}
	if err = tx.Commit(ctx); err != nil {
		return core.Turn{}, false, err
	}
	return stored, false, nil
}

func (s *Store) ResolveSession(ctx context.Context, mapping core.SessionMapping, create func(context.Context) (string, error)) (string, error) {
	tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{})
	if err != nil {
		return "", err
	}
	defer func() { _ = tx.Rollback(ctx) }()
	if _, err = tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1,0))`, mapping.SessionKey); err != nil {
		return "", err
	}
	var sessionID string
	err = tx.QueryRow(ctx, `SELECT weknora_session_id FROM session_mappings WHERE session_key=$1`, mapping.SessionKey).Scan(&sessionID)
	if err == nil {
		return sessionID, tx.Commit(ctx)
	}
	if !errors.Is(err, pgx.ErrNoRows) {
		return "", err
	}
	sessionID, err = create(ctx)
	if err != nil {
		return "", err
	}
	_, err = tx.Exec(ctx, `INSERT INTO session_mappings
		(session_key,tenant_key,channel,agent_id,conversation_type,external_subject_id,weknora_session_id,display_title)
		VALUES($1,$2,$3,$4,$5,$6,$7,$8)`, mapping.SessionKey, mapping.TenantKey, mapping.Channel, mapping.AgentID, string(mapping.ConversationType), mapping.ExternalSubjectID, sessionID, mapping.DisplayTitle)
	if err != nil {
		return "", err
	}
	return sessionID, tx.Commit(ctx)
}

func (s *Store) SyncSessionTitle(ctx context.Context, mapping core.SessionMapping, title string, update core.SessionTitleUpdate) error {
	title = strings.TrimSpace(title)
	if strings.TrimSpace(mapping.SessionKey) == "" || title == "" || update == nil {
		return errors.New("session key, title and update callback are required")
	}
	tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{})
	if err != nil {
		return err
	}
	defer func() { _ = tx.Rollback(ctx) }()
	if _, err = tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1,0))`, mapping.SessionKey); err != nil {
		return err
	}
	var sessionID, currentTitle string
	if err = tx.QueryRow(ctx, `SELECT weknora_session_id,display_title FROM session_mappings WHERE session_key=$1`, mapping.SessionKey).Scan(&sessionID, &currentTitle); err != nil {
		return err
	}
	if currentTitle == title {
		return tx.Commit(ctx)
	}
	if err = update(ctx, sessionID, title); err != nil {
		return fmt.Errorf("update session title: %w", err)
	}
	if _, err = tx.Exec(ctx, `UPDATE session_mappings SET display_title=$2,updated_at=NOW() WHERE session_key=$1`, mapping.SessionKey, title); err != nil {
		return err
	}
	return tx.Commit(ctx)
}

func (s *Store) InvalidateSession(ctx context.Context, sessionKey, expectedSessionID string) error {
	_, err := s.pool.Exec(ctx, `DELETE FROM session_mappings WHERE session_key=$1 AND weknora_session_id=$2`, sessionKey, expectedSessionID)
	return err
}

func (s *Store) FinishTurn(ctx context.Context, turn core.Turn) error {
	references, err := json.Marshal(turn.References)
	if err != nil {
		return err
	}
	media := turn.Media
	if media == nil {
		media = []core.MediaAudit{}
	}
	mediaJSON, err := json.Marshal(media)
	if err != nil {
		return err
	}
	tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{})
	if err != nil {
		return err
	}
	defer func() { _ = tx.Rollback(ctx) }()

	var answerMessageID any
	if shouldCreateAnswerMessage(turn) {
		_, err = tx.Exec(ctx, `INSERT INTO messages
			(id,tenant_key,channel,conversation_id,external_message_id,correlation_id,
			 direction,role,message_type,text_content,media_json,reply_to_message_id,status,actor_name_snapshot,safe_metadata_json)
			VALUES($1,$2,$3,$4,'',$5,'outbound','assistant','text',$6,'[]'::JSONB,$7,'generated','','{}'::JSONB)
			ON CONFLICT (id) DO NOTHING`, turn.AnswerMessageID, turn.TenantKey, turn.Channel,
			turn.ConversationID, turn.ID, turn.Answer, turn.RequestMessageID)
		if err != nil {
			return err
		}
		answerMessageID = turn.AnswerMessageID

		if shouldEnqueueQuestionTableExport(turn, s.questionTableExportEnabled) {
			exportID, err := newUUID()
			if err != nil {
				return err
			}
			_, err = tx.Exec(ctx, `INSERT INTO question_table_exports
				(id,tenant_key,channel,turn_id,status)
				VALUES($1,$2,$3,$4,'pending')
				ON CONFLICT (tenant_key,channel,turn_id) DO NOTHING`,
				exportID, turn.TenantKey, turn.Channel, turn.ID)
			if err != nil {
				return err
			}
		}
	}
	_, err = tx.Exec(ctx, `UPDATE turns SET weknora_session_id=$2,weknora_request_id=$3,answer=$4,references_json=$5,media_json=$6,status=$7,error_message=$8,answer_message_id=COALESCE($9::TEXT,answer_message_id),knowledge_searched=$10,updated_at=NOW() WHERE id=$1`,
		turn.ID, turn.WeKnoraSessionID, turn.WeKnoraRequestID, turn.Answer, references, mediaJSON,
		turn.Status, turn.ErrorMessage, answerMessageID, turn.KnowledgeSearched)
	if err != nil {
		return err
	}
	return tx.Commit(ctx)
}

func (s *Store) RecordDelivery(ctx context.Context, update core.DeliveryUpdate) error {
	if err := validateDeliveryUpdate(update); err != nil {
		return err
	}
	tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{})
	if err != nil {
		return err
	}
	defer func() { _ = tx.Rollback(ctx) }()

	command, err := tx.Exec(ctx, `UPDATE messages AS m SET
		status=$5,
		delivery_method=CASE WHEN $5='sent' THEN $6 ELSE '' END,
		provider_message_id=CASE WHEN $5='sent' THEN $7 ELSE '' END,
		failure_code=CASE WHEN $5='failed' THEN $8 ELSE '' END,
		delivered_at=CASE WHEN $5='sent' THEN COALESCE(m.delivered_at,NOW()) ELSE NULL END,
		updated_at=NOW()
		FROM turns t
		WHERE m.tenant_key=$1 AND m.channel=$2 AND m.id=$3 AND m.correlation_id=$4
		  AND m.direction='outbound' AND m.role='assistant'
		  AND t.id=$4 AND t.tenant_key=$1 AND t.channel=$2 AND t.answer_message_id=m.id
		  AND (($5='sent' AND m.status IN ('generated','failed','sent'))
		       OR ($5='failed' AND (m.status='generated' OR m.status='failed')))`,
		update.TenantKey, update.Channel, update.AnswerMessageID, update.CorrelationID,
		update.Status, update.DeliveryMethod, update.ProviderMessageID, update.FailureCode)
	if err != nil {
		return err
	}
	if command.RowsAffected() != 1 {
		return core.ErrDeliveryUpdateFailed
	}
	return tx.Commit(ctx)
}

type rowQuerier interface {
	QueryRow(context.Context, string, ...any) pgx.Row
}

func (s *Store) loadTurn(ctx context.Context, tenantKey, channel, messageID string) (core.Turn, error) {
	return loadTurnWith(ctx, s.pool, tenantKey, channel, messageID)
}

func loadTurnWith(ctx context.Context, query rowQuerier, tenantKey, channel, messageID string) (core.Turn, error) {
	var turn core.Turn
	var references []byte
	var mediaJSON []byte
	var agentID, requestMessageID, answerMessageID *string
	var actorID, conversationID, messageType, actorName, actorUnionID *string
	var conversationType, externalConversationID, externalActorID *string
	var answerDeliveryStatus, answerDeliveryMethod *string
	err := query.QueryRow(ctx, `SELECT
		t.id,t.tenant_key,t.channel,t.external_message_id,t.session_key,t.weknora_session_id,
		t.weknora_request_id,t.question,t.answer,t.references_json,t.media_json,t.status,t.error_message,
		t.agent_id,t.request_message_id,t.answer_message_id,
		m.actor_id,m.conversation_id,m.message_type,m.actor_name_snapshot,m.actor_union_id_snapshot,
		c.conversation_type,c.external_conversation_id,a.external_actor_id,
		am.status,am.delivery_method
		FROM turns t
		LEFT JOIN messages m ON m.id=t.request_message_id
		LEFT JOIN messages am ON am.id=t.answer_message_id
		LEFT JOIN channel_conversations c ON c.id=m.conversation_id
		LEFT JOIN channel_actors a ON a.id=m.actor_id
		WHERE t.tenant_key=$1 AND t.channel=$2 AND t.external_message_id=$3`, tenantKey, channel, messageID).Scan(
		&turn.ID, &turn.TenantKey, &turn.Channel, &turn.ExternalMessageID, &turn.SessionKey,
		&turn.WeKnoraSessionID, &turn.WeKnoraRequestID, &turn.Question, &turn.Answer, &references,
		&mediaJSON, &turn.Status, &turn.ErrorMessage, &agentID, &requestMessageID, &answerMessageID,
		&actorID, &conversationID, &messageType, &actorName, &actorUnionID, &conversationType,
		&externalConversationID, &externalActorID, &answerDeliveryStatus, &answerDeliveryMethod)
	if err != nil {
		return core.Turn{}, err
	}
	if agentID != nil {
		turn.AgentID = *agentID
	}
	if requestMessageID != nil {
		turn.RequestMessageID = *requestMessageID
	}
	if answerMessageID != nil {
		turn.AnswerMessageID = *answerMessageID
	}
	if answerDeliveryStatus != nil {
		turn.AnswerDeliveryStatus = *answerDeliveryStatus
	}
	if answerDeliveryMethod != nil {
		turn.AnswerDeliveryMethod = *answerDeliveryMethod
	}
	if actorID != nil {
		turn.ActorID = *actorID
	}
	if conversationID != nil {
		turn.ConversationID = *conversationID
	}
	if messageType != nil {
		turn.MessageType = *messageType
	}
	if actorName != nil {
		turn.SenderName = *actorName
	}
	if actorUnionID != nil {
		turn.ActorUnionIDSnapshot = *actorUnionID
	}
	if externalActorID != nil {
		turn.ExternalUserID = *externalActorID
	}
	if conversationType != nil {
		turn.ConversationType = core.ConversationType(*conversationType)
	}
	if externalConversationID != nil && (turn.ConversationType == core.ConversationGroup ||
		(turn.ExternalUserID != "" && *externalConversationID != turn.ExternalUserID)) {
		turn.ExternalChatID = *externalConversationID
	}
	if len(references) > 0 {
		if err := json.Unmarshal(references, &turn.References); err != nil {
			return core.Turn{}, err
		}
	}
	if len(mediaJSON) > 0 {
		if err := json.Unmarshal(mediaJSON, &turn.Media); err != nil {
			return core.Turn{}, err
		}
	}
	return turn, nil
}

func shouldCreateAnswerMessage(turn core.Turn) bool {
	return turn.Status == "succeeded" && strings.TrimSpace(turn.Answer) != "" &&
		strings.TrimSpace(turn.AnswerMessageID) != ""
}

func shouldEnqueueQuestionTableExport(turn core.Turn, enabled bool) bool {
	return enabled && shouldCreateAnswerMessage(turn) && strings.EqualFold(strings.TrimSpace(turn.Channel), "dingtalk")
}

func validateDeliveryUpdate(update core.DeliveryUpdate) error {
	if strings.TrimSpace(update.TenantKey) == "" || strings.TrimSpace(update.Channel) == "" ||
		strings.TrimSpace(update.AnswerMessageID) == "" || strings.TrimSpace(update.CorrelationID) == "" {
		return errors.New("delivery update identity is required")
	}
	switch update.Status {
	case core.DeliverySent:
		if update.DeliveryMethod != core.DeliveryMethodCard && update.DeliveryMethod != core.DeliveryMethodMarkdown {
			return errors.New("sent delivery requires a supported method")
		}
		if update.FailureCode != "" {
			return errors.New("sent delivery cannot have a failure code")
		}
		if update.DeliveryMethod == core.DeliveryMethodMarkdown && update.ProviderMessageID != "" {
			return errors.New("markdown delivery cannot have a provider message id")
		}
	case core.DeliveryFailed:
		if update.DeliveryMethod != "" || update.ProviderMessageID != "" ||
			update.FailureCode != core.DeliveryFailureChannelDelivery {
			return errors.New("failed delivery requires a safe failure code only")
		}
	default:
		return errors.New("delivery status must be sent or failed")
	}
	return nil
}

func externalConversationID(turn core.Turn) string {
	if externalChatID := strings.TrimSpace(turn.ExternalChatID); externalChatID != "" {
		return externalChatID
	}
	return strings.TrimSpace(turn.ExternalUserID)
}