gooseek/backend/internal/db/computer_artifact_repo.go

package db

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

	"github.com/gooseek/backend/internal/computer"
)

type ComputerArtifactRepo struct {
	db *sql.DB
}

func NewComputerArtifactRepo(db *sql.DB) *ComputerArtifactRepo {
	return &ComputerArtifactRepo{db: db}
}

func (r *ComputerArtifactRepo) Migrate() error {
	query := `
	CREATE TABLE IF NOT EXISTS computer_artifacts (
		id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
		task_id UUID NOT NULL,
		type VARCHAR(50) NOT NULL,
		name VARCHAR(255),
		content BYTEA,
		url TEXT,
		size BIGINT DEFAULT 0,
		mime_type VARCHAR(100),
		metadata JSONB,
		created_at TIMESTAMPTZ DEFAULT NOW()
	);

	CREATE INDEX IF NOT EXISTS idx_computer_artifacts_task_id ON computer_artifacts(task_id);
	CREATE INDEX IF NOT EXISTS idx_computer_artifacts_type ON computer_artifacts(type);
	CREATE INDEX IF NOT EXISTS idx_computer_artifacts_created ON computer_artifacts(created_at DESC);
	`

	_, err := r.db.Exec(query)
	return err
}

func (r *ComputerArtifactRepo) Create(ctx context.Context, artifact *computer.Artifact) error {
	metadataJSON, _ := json.Marshal(artifact.Metadata)

	query := `
		INSERT INTO computer_artifacts (id, task_id, type, name, content, url, size, mime_type, metadata, created_at)
		VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
	`

	_, err := r.db.ExecContext(ctx, query,
		artifact.ID,
		artifact.TaskID,
		artifact.Type,
		artifact.Name,
		artifact.Content,
		artifact.URL,
		artifact.Size,
		artifact.MimeType,
		metadataJSON,
		artifact.CreatedAt,
	)

	return err
}

func (r *ComputerArtifactRepo) GetByID(ctx context.Context, id string) (*computer.Artifact, error) {
	query := `
		SELECT id, task_id, type, name, content, url, size, mime_type, metadata, created_at
		FROM computer_artifacts
		WHERE id = $1
	`

	var artifact computer.Artifact
	var content []byte
	var url, mimeType sql.NullString
	var metadataJSON []byte

	err := r.db.QueryRowContext(ctx, query, id).Scan(
		&artifact.ID,
		&artifact.TaskID,
		&artifact.Type,
		&artifact.Name,
		&content,
		&url,
		&artifact.Size,
		&mimeType,
		&metadataJSON,
		&artifact.CreatedAt,
	)

	if err != nil {
		return nil, err
	}

	artifact.Content = content
	if url.Valid {
		artifact.URL = url.String
	}
	if mimeType.Valid {
		artifact.MimeType = mimeType.String
	}
	if len(metadataJSON) > 0 {
		json.Unmarshal(metadataJSON, &artifact.Metadata)
	}

	return &artifact, nil
}

func (r *ComputerArtifactRepo) GetByTaskID(ctx context.Context, taskID string) ([]computer.Artifact, error) {
	query := `
		SELECT id, task_id, type, name, url, size, mime_type, metadata, created_at
		FROM computer_artifacts
		WHERE task_id = $1
		ORDER BY created_at ASC
	`

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

	var artifacts []computer.Artifact

	for rows.Next() {
		var artifact computer.Artifact
		var url, mimeType sql.NullString
		var metadataJSON []byte

		err := rows.Scan(
			&artifact.ID,
			&artifact.TaskID,
			&artifact.Type,
			&artifact.Name,
			&url,
			&artifact.Size,
			&mimeType,
			&metadataJSON,
			&artifact.CreatedAt,
		)
		if err != nil {
			continue
		}

		if url.Valid {
			artifact.URL = url.String
		}
		if mimeType.Valid {
			artifact.MimeType = mimeType.String
		}
		if len(metadataJSON) > 0 {
			json.Unmarshal(metadataJSON, &artifact.Metadata)
		}

		artifacts = append(artifacts, artifact)
	}

	return artifacts, nil
}

func (r *ComputerArtifactRepo) GetByType(ctx context.Context, taskID, artifactType string) ([]computer.Artifact, error) {
	query := `
		SELECT id, task_id, type, name, url, size, mime_type, metadata, created_at
		FROM computer_artifacts
		WHERE task_id = $1 AND type = $2
		ORDER BY created_at ASC
	`

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

	var artifacts []computer.Artifact

	for rows.Next() {
		var artifact computer.Artifact
		var url, mimeType sql.NullString
		var metadataJSON []byte

		err := rows.Scan(
			&artifact.ID,
			&artifact.TaskID,
			&artifact.Type,
			&artifact.Name,
			&url,
			&artifact.Size,
			&mimeType,
			&metadataJSON,
			&artifact.CreatedAt,
		)
		if err != nil {
			continue
		}

		if url.Valid {
			artifact.URL = url.String
		}
		if mimeType.Valid {
			artifact.MimeType = mimeType.String
		}
		if len(metadataJSON) > 0 {
			json.Unmarshal(metadataJSON, &artifact.Metadata)
		}

		artifacts = append(artifacts, artifact)
	}

	return artifacts, nil
}

func (r *ComputerArtifactRepo) GetContent(ctx context.Context, id string) ([]byte, error) {
	query := `SELECT content FROM computer_artifacts WHERE id = $1`
	var content []byte
	err := r.db.QueryRowContext(ctx, query, id).Scan(&content)
	return content, err
}

func (r *ComputerArtifactRepo) UpdateURL(ctx context.Context, id, url string) error {
	query := `UPDATE computer_artifacts SET url = $1 WHERE id = $2`
	_, err := r.db.ExecContext(ctx, query, url, id)
	return err
}

func (r *ComputerArtifactRepo) Delete(ctx context.Context, id string) error {
	query := `DELETE FROM computer_artifacts WHERE id = $1`
	_, err := r.db.ExecContext(ctx, query, id)
	return err
}

func (r *ComputerArtifactRepo) DeleteByTaskID(ctx context.Context, taskID string) error {
	query := `DELETE FROM computer_artifacts WHERE task_id = $1`
	_, err := r.db.ExecContext(ctx, query, taskID)
	return err
}

func (r *ComputerArtifactRepo) DeleteOlderThan(ctx context.Context, days int) (int64, error) {
	query := `
		DELETE FROM computer_artifacts
		WHERE created_at < NOW() - INTERVAL '1 day' * $1
	`
	result, err := r.db.ExecContext(ctx, query, days)
	if err != nil {
		return 0, err
	}
	return result.RowsAffected()
}

func (r *ComputerArtifactRepo) GetTotalSize(ctx context.Context, taskID string) (int64, error) {
	query := `SELECT COALESCE(SUM(size), 0) FROM computer_artifacts WHERE task_id = $1`
	var size int64
	err := r.db.QueryRowContext(ctx, query, taskID).Scan(&size)
	return size, err
}

func (r *ComputerArtifactRepo) Count(ctx context.Context, taskID string) (int64, error) {
	query := `SELECT COUNT(*) FROM computer_artifacts WHERE task_id = $1`
	var count int64
	err := r.db.QueryRowContext(ctx, query, taskID).Scan(&count)
	return count, err
}

type ArtifactSummary struct {
	ID        string    `json:"id"`
	TaskID    string    `json:"taskId"`
	Type      string    `json:"type"`
	Name      string    `json:"name"`
	URL       string    `json:"url"`
	Size      int64     `json:"size"`
	MimeType  string    `json:"mimeType"`
	CreatedAt time.Time `json:"createdAt"`
}

func (r *ComputerArtifactRepo) GetSummaries(ctx context.Context, taskID string) ([]ArtifactSummary, error) {
	query := `
		SELECT id, task_id, type, name, url, size, mime_type, created_at
		FROM computer_artifacts
		WHERE task_id = $1
		ORDER BY created_at ASC
	`

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

	var summaries []ArtifactSummary

	for rows.Next() {
		var s ArtifactSummary
		var url, mimeType sql.NullString

		err := rows.Scan(
			&s.ID,
			&s.TaskID,
			&s.Type,
			&s.Name,
			&url,
			&s.Size,
			&mimeType,
			&s.CreatedAt,
		)
		if err != nil {
			continue
		}

		if url.Valid {
			s.URL = url.String
		}
		if mimeType.Valid {
			s.MimeType = mimeType.String
		}

		summaries = append(summaries, s)
	}

	return summaries, nil
}