package repository

import (
	"context"
	"errors"
	"time"

	"github.com/jackc/pgx/v5"

	"github.com/geo-platform/tenant-api/internal/tenant/repository/generated"
)

type KolAssistTask struct {
	ID           string
	TenantID     int64
	KolProfileID int64
	PromptID     *int64
	OperatorID   int64
	TaskType     string
	Status       string
	RequestJSON  []byte
	ResultJSON   []byte
	ErrorMessage *string
	StartedAt    *time.Time
	CompletedAt  *time.Time
	CreatedAt    time.Time
	UpdatedAt    time.Time
}

type CreateKolAssistInput struct {
	ID           string
	TenantID     int64
	KolProfileID int64
	PromptID     *int64
	OperatorID   int64
	TaskType     string
	RequestJSON  []byte
}

type KolAssistRepository interface {
	Create(ctx context.Context, input CreateKolAssistInput) error
	Get(ctx context.Context, tenantID int64, id string) (*KolAssistTask, error)
	MarkStarted(ctx context.Context, tenantID int64, id string) error
	MarkCompleted(ctx context.Context, tenantID int64, id string, resultJSON []byte) error
	MarkFailed(ctx context.Context, tenantID int64, id string, errorMsg string) error
}

type kolAssistRepository struct {
	q generated.Querier
}

func NewKolAssistRepository(db generated.DBTX) KolAssistRepository {
	return &kolAssistRepository{q: newQuerier(db)}
}

func (r *kolAssistRepository) Create(ctx context.Context, input CreateKolAssistInput) error {
	return r.q.CreateKolAssistTask(ctx, generated.CreateKolAssistTaskParams{
		ID:           input.ID,
		TenantID:     input.TenantID,
		KolProfileID: input.KolProfileID,
		PromptID:     pgInt8(input.PromptID),
		OperatorID:   input.OperatorID,
		TaskType:     input.TaskType,
		RequestJson:  input.RequestJSON,
	})
}

func (r *kolAssistRepository) Get(ctx context.Context, tenantID int64, id string) (*KolAssistTask, error) {
	row, err := r.q.GetKolAssistTask(ctx, generated.GetKolAssistTaskParams{
		ID:       id,
		TenantID: tenantID,
	})
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, nil
		}
		return nil, err
	}
	return &KolAssistTask{
		ID:           row.ID,
		TenantID:     row.TenantID,
		KolProfileID: row.KolProfileID,
		PromptID:     nullableInt64(row.PromptID),
		OperatorID:   row.OperatorID,
		TaskType:     row.TaskType,
		Status:       row.Status,
		RequestJSON:  row.RequestJson,
		ResultJSON:   row.ResultJson,
		ErrorMessage: nullableText(row.ErrorMessage),
		StartedAt:    optionalTime(row.StartedAt),
		CompletedAt:  optionalTime(row.CompletedAt),
		CreatedAt:    timeFromTimestamp(row.CreatedAt),
		UpdatedAt:    timeFromTimestamp(row.UpdatedAt),
	}, nil
}

func (r *kolAssistRepository) MarkStarted(ctx context.Context, tenantID int64, id string) error {
	return r.q.MarkKolAssistStarted(ctx, generated.MarkKolAssistStartedParams{
		ID:       id,
		TenantID: tenantID,
	})
}

func (r *kolAssistRepository) MarkCompleted(ctx context.Context, tenantID int64, id string, resultJSON []byte) error {
	return r.q.MarkKolAssistCompleted(ctx, generated.MarkKolAssistCompletedParams{
		ResultJson: resultJSON,
		ID:         id,
		TenantID:   tenantID,
	})
}

func (r *kolAssistRepository) MarkFailed(ctx context.Context, tenantID int64, id string, errorMsg string) error {
	return r.q.MarkKolAssistFailed(ctx, generated.MarkKolAssistFailedParams{
		ErrorMessage: pgText(&errorMsg),
		ID:           id,
		TenantID:     tenantID,
	})
}