gooseek/backend/internal/db/thread_repo.go

package db

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

var (
	ErrNotFound  = errors.New("resource not found")
	ErrForbidden = errors.New("access denied")
)

type Thread struct {
	ID          string          `json:"id"`
	UserID      string          `json:"userId"`
	SpaceID     *string         `json:"spaceId,omitempty"`
	Title       string          `json:"title"`
	FocusMode   string          `json:"focusMode"`
	IsPublic    bool            `json:"isPublic"`
	ShareID     *string         `json:"shareId,omitempty"`
	CreatedAt   time.Time       `json:"createdAt"`
	UpdatedAt   time.Time       `json:"updatedAt"`
	Messages    []ThreadMessage `json:"messages,omitempty"`
	MessageCount int            `json:"messageCount,omitempty"`
}

type ThreadMessage struct {
	ID            string                 `json:"id"`
	ThreadID      string                 `json:"threadId"`
	Role          string                 `json:"role"`
	Content       string                 `json:"content"`
	Sources       []ThreadSource         `json:"sources,omitempty"`
	Widgets       []map[string]interface{} `json:"widgets,omitempty"`
	RelatedQuestions []string            `json:"relatedQuestions,omitempty"`
	Model         string                 `json:"model,omitempty"`
	TokensUsed    int                    `json:"tokensUsed,omitempty"`
	CreatedAt     time.Time              `json:"createdAt"`
}

type ThreadSource struct {
	Index   int    `json:"index"`
	URL     string `json:"url"`
	Title   string `json:"title"`
	Domain  string `json:"domain"`
	Snippet string `json:"snippet,omitempty"`
}

type ThreadRepository struct {
	db *PostgresDB
}

func NewThreadRepository(db *PostgresDB) *ThreadRepository {
	return &ThreadRepository{db: db}
}

func (r *ThreadRepository) RunMigrations(ctx context.Context) error {
	migrations := []string{
		`CREATE TABLE IF NOT EXISTS threads (
			id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
			user_id UUID NOT NULL,
			space_id UUID REFERENCES spaces(id) ON DELETE SET NULL,
			title VARCHAR(500) NOT NULL DEFAULT 'New Thread',
			focus_mode VARCHAR(50) DEFAULT 'all',
			is_public BOOLEAN DEFAULT FALSE,
			share_id VARCHAR(100) UNIQUE,
			created_at TIMESTAMPTZ DEFAULT NOW(),
			updated_at TIMESTAMPTZ DEFAULT NOW()
		)`,
		`CREATE INDEX IF NOT EXISTS idx_threads_user ON threads(user_id)`,
		`CREATE INDEX IF NOT EXISTS idx_threads_space ON threads(space_id)`,
		`CREATE INDEX IF NOT EXISTS idx_threads_share ON threads(share_id)`,
		`CREATE TABLE IF NOT EXISTS thread_messages (
			id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
			thread_id UUID NOT NULL REFERENCES threads(id) ON DELETE CASCADE,
			role VARCHAR(20) NOT NULL,
			content TEXT NOT NULL,
			sources JSONB DEFAULT '[]',
			widgets JSONB DEFAULT '[]',
			related_questions JSONB DEFAULT '[]',
			model VARCHAR(100),
			tokens_used INT DEFAULT 0,
			created_at TIMESTAMPTZ DEFAULT NOW()
		)`,
		`CREATE INDEX IF NOT EXISTS idx_thread_messages_thread ON thread_messages(thread_id)`,
	}

	for _, m := range migrations {
		if _, err := r.db.db.ExecContext(ctx, m); err != nil {
			return err
		}
	}
	return nil
}

func (r *ThreadRepository) Create(ctx context.Context, t *Thread) error {
	query := `
		INSERT INTO threads (user_id, space_id, title, focus_mode, is_public)
		VALUES ($1, $2, $3, $4, $5)
		RETURNING id, created_at, updated_at
	`
	return r.db.db.QueryRowContext(ctx, query,
		t.UserID, t.SpaceID, t.Title, t.FocusMode, t.IsPublic,
	).Scan(&t.ID, &t.CreatedAt, &t.UpdatedAt)
}

func (r *ThreadRepository) GetByID(ctx context.Context, id string) (*Thread, error) {
	query := `
		SELECT id, user_id, space_id, title, focus_mode, is_public, share_id, created_at, updated_at,
			   (SELECT COUNT(*) FROM thread_messages WHERE thread_id = threads.id) as message_count
		FROM threads 
		WHERE id = $1
	`

	var t Thread
	err := r.db.db.QueryRowContext(ctx, query, id).Scan(
		&t.ID, &t.UserID, &t.SpaceID, &t.Title, &t.FocusMode,
		&t.IsPublic, &t.ShareID, &t.CreatedAt, &t.UpdatedAt, &t.MessageCount,
	)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}

	return &t, nil
}

func (r *ThreadRepository) GetByShareID(ctx context.Context, shareID string) (*Thread, error) {
	query := `
		SELECT id, user_id, space_id, title, focus_mode, is_public, share_id, created_at, updated_at
		FROM threads 
		WHERE share_id = $1 AND is_public = true
	`

	var t Thread
	err := r.db.db.QueryRowContext(ctx, query, shareID).Scan(
		&t.ID, &t.UserID, &t.SpaceID, &t.Title, &t.FocusMode,
		&t.IsPublic, &t.ShareID, &t.CreatedAt, &t.UpdatedAt,
	)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}

	return &t, nil
}

func (r *ThreadRepository) GetByUserID(ctx context.Context, userID string, limit, offset int) ([]*Thread, error) {
	query := `
		SELECT id, user_id, space_id, title, focus_mode, is_public, share_id, created_at, updated_at,
			   (SELECT COUNT(*) FROM thread_messages WHERE thread_id = threads.id) as message_count
		FROM threads 
		WHERE user_id = $1
		ORDER BY updated_at DESC
		LIMIT $2 OFFSET $3
	`

	rows, err := r.db.db.QueryContext(ctx, query, userID, limit, offset)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var threads []*Thread
	for rows.Next() {
		var t Thread
		if err := rows.Scan(
			&t.ID, &t.UserID, &t.SpaceID, &t.Title, &t.FocusMode,
			&t.IsPublic, &t.ShareID, &t.CreatedAt, &t.UpdatedAt, &t.MessageCount,
		); err != nil {
			return nil, err
		}
		threads = append(threads, &t)
	}

	return threads, nil
}

func (r *ThreadRepository) Update(ctx context.Context, t *Thread, userID string) error {
	query := `
		UPDATE threads 
		SET title = $2, focus_mode = $3, is_public = $4, updated_at = NOW()
		WHERE id = $1 AND user_id = $5
	`
	result, err := r.db.db.ExecContext(ctx, query, t.ID, t.Title, t.FocusMode, t.IsPublic, userID)
	if err != nil {
		return err
	}
	rows, _ := result.RowsAffected()
	if rows == 0 {
		return ErrNotFound
	}
	return nil
}

func (r *ThreadRepository) SetShareID(ctx context.Context, threadID, shareID, userID string) error {
	result, err := r.db.db.ExecContext(ctx,
		"UPDATE threads SET share_id = $2, is_public = true WHERE id = $1 AND user_id = $3",
		threadID, shareID, userID,
	)
	if err != nil {
		return err
	}
	rows, _ := result.RowsAffected()
	if rows == 0 {
		return ErrNotFound
	}
	return nil
}

func (r *ThreadRepository) Delete(ctx context.Context, id, userID string) error {
	result, err := r.db.db.ExecContext(ctx, "DELETE FROM threads WHERE id = $1 AND user_id = $2", id, userID)
	if err != nil {
		return err
	}
	rows, _ := result.RowsAffected()
	if rows == 0 {
		return ErrNotFound
	}
	return nil
}

func (r *ThreadRepository) AddMessage(ctx context.Context, msg *ThreadMessage, userID string) error {
	var ownerID string
	err := r.db.db.QueryRowContext(ctx, "SELECT user_id FROM threads WHERE id = $1", msg.ThreadID).Scan(&ownerID)
	if err == sql.ErrNoRows {
		return ErrNotFound
	}
	if err != nil {
		return err
	}
	if ownerID != userID {
		return ErrForbidden
	}

	sourcesJSON, _ := json.Marshal(msg.Sources)
	widgetsJSON, _ := json.Marshal(msg.Widgets)
	relatedJSON, _ := json.Marshal(msg.RelatedQuestions)

	query := `
		INSERT INTO thread_messages (thread_id, role, content, sources, widgets, related_questions, model, tokens_used)
		VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
		RETURNING id, created_at
	`
	err = r.db.db.QueryRowContext(ctx, query,
		msg.ThreadID, msg.Role, msg.Content, sourcesJSON, widgetsJSON, relatedJSON, msg.Model, msg.TokensUsed,
	).Scan(&msg.ID, &msg.CreatedAt)

	if err == nil {
		r.db.db.ExecContext(ctx, "UPDATE threads SET updated_at = NOW() WHERE id = $1", msg.ThreadID)
	}

	return err
}

func (r *ThreadRepository) GetMessages(ctx context.Context, threadID, userID string, limit, offset int) ([]ThreadMessage, error) {
	var ownerID string
	var isPublic bool
	err := r.db.db.QueryRowContext(ctx, "SELECT user_id, is_public FROM threads WHERE id = $1", threadID).Scan(&ownerID, &isPublic)
	if err == sql.ErrNoRows {
		return nil, ErrNotFound
	}
	if err != nil {
		return nil, err
	}
	if ownerID != userID && !isPublic {
		return nil, ErrForbidden
	}

	query := `
		SELECT id, thread_id, role, content, sources, widgets, related_questions, model, tokens_used, created_at
		FROM thread_messages 
		WHERE thread_id = $1
		ORDER BY created_at ASC
		LIMIT $2 OFFSET $3
	`

	rows, err := r.db.db.QueryContext(ctx, query, threadID, limit, offset)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var messages []ThreadMessage
	for rows.Next() {
		var msg ThreadMessage
		var sourcesJSON, widgetsJSON, relatedJSON []byte

		if err := rows.Scan(
			&msg.ID, &msg.ThreadID, &msg.Role, &msg.Content,
			&sourcesJSON, &widgetsJSON, &relatedJSON,
			&msg.Model, &msg.TokensUsed, &msg.CreatedAt,
		); err != nil {
			return nil, err
		}

		json.Unmarshal(sourcesJSON, &msg.Sources)
		json.Unmarshal(widgetsJSON, &msg.Widgets)
		json.Unmarshal(relatedJSON, &msg.RelatedQuestions)
		messages = append(messages, msg)
	}

	return messages, nil
}

func (r *ThreadRepository) GenerateTitle(ctx context.Context, threadID, firstMessage, userID string) error {
	title := firstMessage
	if len(title) > 100 {
		title = title[:97] + "..."
	}
	result, err := r.db.db.ExecContext(ctx,
		"UPDATE threads SET title = $2 WHERE id = $1 AND user_id = $3",
		threadID, title, userID,
	)
	if err != nil {
		return err
	}
	rows, _ := result.RowsAffected()
	if rows == 0 {
		return ErrNotFound
	}
	return nil
}