package repository

import (
	"context"
	"errors"
	"time"

	"github.com/jackc/pgx/v5"

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

var ErrTemplateAssistTaskNotFound = errors.New("template assist task not found")

type TemplateAssistTaskRecord struct {
	ID           string
	Status       string
	ResultJSON   []byte
	ErrorMessage *string
}

type CreateTemplateAssistTaskInput struct {
	ID          string
	TaskType    string
	TenantID    int64
	OperatorID  int64
	TemplateID  int64
	RequestJSON []byte
}

type TemplateAssistTaskRepository interface {
	CreateTask(ctx context.Context, input CreateTemplateAssistTaskInput) error
	GetTask(ctx context.Context, id, taskType string, tenantID, templateID int64) (*TemplateAssistTaskRecord, error)
	MarkRunning(ctx context.Context, id, taskType string, tenantID int64, startedAt time.Time) error
	CompleteTask(ctx context.Context, id, taskType string, tenantID int64, resultJSON []byte, completedAt time.Time) error
	FailTask(ctx context.Context, id, taskType string, tenantID int64, errorMessage string, completedAt time.Time) error
}

type templateAssistTaskRepository struct {
	db generated.DBTX
}

func NewTemplateAssistTaskRepository(db generated.DBTX) TemplateAssistTaskRepository {
	return &templateAssistTaskRepository{db: db}
}

func (r *templateAssistTaskRepository) CreateTask(ctx context.Context, input CreateTemplateAssistTaskInput) error {
	_, err := r.db.Exec(ctx, `
		INSERT INTO template_assist_tasks (id, task_type, tenant_id, operator_id, template_id, request_json, status)
		VALUES ($1, $2, $3, $4, $5, $6, 'queued')
	`, input.ID, input.TaskType, input.TenantID, input.OperatorID, input.TemplateID, input.RequestJSON)
	return err
}

func (r *templateAssistTaskRepository) GetTask(ctx context.Context, id, taskType string, tenantID, templateID int64) (*TemplateAssistTaskRecord, error) {
	var record TemplateAssistTaskRecord
	err := r.db.QueryRow(ctx, `
		SELECT id, status, result_json, error_message
		FROM template_assist_tasks
		WHERE id = $1 AND task_type = $2 AND tenant_id = $3 AND template_id = $4
	`, id, taskType, tenantID, templateID).Scan(&record.ID, &record.Status, &record.ResultJSON, &record.ErrorMessage)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, ErrTemplateAssistTaskNotFound
		}
		return nil, err
	}
	return &record, nil
}

func (r *templateAssistTaskRepository) MarkRunning(ctx context.Context, id, taskType string, tenantID int64, startedAt time.Time) error {
	_, err := r.db.Exec(ctx, `
		UPDATE template_assist_tasks
		SET status = 'running', started_at = $1, updated_at = NOW()
		WHERE id = $2 AND task_type = $3 AND tenant_id = $4
	`, startedAt, id, taskType, tenantID)
	return err
}

func (r *templateAssistTaskRepository) CompleteTask(ctx context.Context, id, taskType string, tenantID int64, resultJSON []byte, completedAt time.Time) error {
	_, err := r.db.Exec(ctx, `
		UPDATE template_assist_tasks
		SET status = 'completed',
		    result_json = $1,
		    error_message = NULL,
		    completed_at = $2,
		    updated_at = NOW()
		WHERE id = $3 AND task_type = $4 AND tenant_id = $5
	`, resultJSON, completedAt, id, taskType, tenantID)
	return err
}

func (r *templateAssistTaskRepository) FailTask(ctx context.Context, id, taskType string, tenantID int64, errorMessage string, completedAt time.Time) error {
	_, err := r.db.Exec(ctx, `
		UPDATE template_assist_tasks
		SET status = 'failed',
		    error_message = $1,
		    completed_at = $2,
		    updated_at = NOW()
		WHERE id = $3 AND task_type = $4 AND tenant_id = $5
	`, errorMessage, completedAt, id, taskType, tenantID)
	return err
}