package task

import (
	"context"
	"encoding/csv"
	"errors"
	"fmt"
	"io"
	"os"

	"github.com/hibiken/asynq"
	"github.com/redis/go-redis/v9"
	"go.uber.org/zap"

	"github.com/break/junhong_cmp_fiber/internal/exporter"
	"github.com/break/junhong_cmp_fiber/internal/model"
	"github.com/break/junhong_cmp_fiber/internal/store/postgres"
	"github.com/break/junhong_cmp_fiber/pkg/constants"
	"github.com/break/junhong_cmp_fiber/pkg/storage"
)

// ExportFinalizeHandler 导出 finalize 处理器。
type ExportFinalizeHandler struct {
	redis         *redis.Client
	taskStore     *postgres.ExportTaskStore
	shardStore    *postgres.ExportShardTaskStore
	storageSvc    *storage.Service
	sceneRegistry *exporter.Registry
	logger        *zap.Logger
}

// NewExportFinalizeHandler 创建导出 finalize 处理器。
func NewExportFinalizeHandler(
	redis *redis.Client,
	taskStore *postgres.ExportTaskStore,
	shardStore *postgres.ExportShardTaskStore,
	storageSvc *storage.Service,
	sceneRegistry *exporter.Registry,
	logger *zap.Logger,
) *ExportFinalizeHandler {
	return &ExportFinalizeHandler{
		redis:         redis,
		taskStore:     taskStore,
		shardStore:    shardStore,
		storageSvc:    storageSvc,
		sceneRegistry: sceneRegistry,
		logger:        logger,
	}
}

// HandleExportFinalize 处理导出 finalize 任务。
func (h *ExportFinalizeHandler) HandleExportFinalize(ctx context.Context, task *asynq.Task) error {
	var payload ExportFinalizePayload
	if err := parseExportPayload(task.Payload(), &payload); err != nil {
		h.logger.Error("解析导出 finalize 任务载荷失败", zap.Error(err))
		return asynq.SkipRetry
	}

	lockKey := constants.RedisExportFinalizeLockKey(payload.TaskID)
	locked, err := tryAcquireLock(ctx, h.redis, lockKey, exportFinalizeLockTTL)
	if err != nil {
		return err
	}
	if !locked {
		return nil
	}
	defer releaseLock(ctx, h.redis, lockKey)

	exportTask, err := h.taskStore.GetByIDForWorker(ctx, payload.TaskID)
	if err != nil {
		h.logger.Error("查询导出任务失败", zap.Uint("task_id", payload.TaskID), zap.Error(err))
		return asynq.SkipRetry
	}

	if exportTask.Status == constants.ExportTaskStatusCompleted || exportTask.Status == constants.ExportTaskStatusFailed || exportTask.Status == constants.ExportTaskStatusCancelled {
		return nil
	}

	updater := dispatchUpdater(exportTask)
	if exportTask.CancelRequested {
		if err := h.shardStore.CancelByTaskID(ctx, exportTask.ID, updater, "任务已取消"); err != nil {
			return err
		}
		if err := h.taskStore.MarkCancelledByWorker(ctx, exportTask.ID, updater, "任务已取消"); err != nil {
			return err
		}
		return nil
	}

	unfinishedCount, err := h.shardStore.CountUnfinishedByTaskID(ctx, exportTask.ID)
	if err != nil {
		return err
	}
	if unfinishedCount > 0 {
		return fmt.Errorf("分片任务尚未完成: %d", unfinishedCount)
	}

	failedCount, err := h.shardStore.CountByTaskIDAndStatus(ctx, exportTask.ID, constants.ExportShardStatusFailed)
	if err != nil {
		return err
	}
	if failedCount > 0 {
		if err := h.taskStore.MarkFailed(ctx, exportTask.ID, updater, "存在失败分片，导出任务失败"); err != nil {
			return err
		}
		return nil
	}

	if !h.sceneRegistry.IsSupported(exportTask.Scene) {
		if err := h.taskStore.MarkFailed(ctx, exportTask.ID, updater, "导出场景未注册"); err != nil {
			return err
		}
		return nil
	}

	shards, err := h.shardStore.ListByTaskID(ctx, exportTask.ID)
	if err != nil {
		return err
	}

	headers := exporter.ResolvedHeadersFromTask(exportTask)
	if len(headers) == 0 {
		return h.handleFinalizeError(ctx, exportTask.ID, updater, "读取导出表头失败", errors.New("resolved_headers 为空"))
	}

	mergedCSVPath, cleanupCSV, err := h.mergeShardCSVFiles(ctx, headers, shards)
	if err != nil {
		return h.handleFinalizeError(ctx, exportTask.ID, updater, "生成汇总文件失败", err)
	}
	defer cleanupCSV()

	localPath, fileSize, cleanupFinal, err := prepareFinalExportFile(exportTask.Format, mergedCSVPath)
	if err != nil {
		return h.handleFinalizeError(ctx, exportTask.ID, updater, "生成汇总文件失败", err)
	}
	defer cleanupFinal()

	fileName := fmt.Sprintf("%s.%s", exportTask.TaskNo, exportTask.Format)
	fileKey, err := uploadExportFile(ctx, h.storageSvc, localPath, fileName)
	if err != nil {
		return h.handleFinalizeError(ctx, exportTask.ID, updater, "上传汇总文件失败", err)
	}

	if err := h.taskStore.MarkCompleted(ctx, exportTask.ID, updater, fileKey, fileSize); err != nil {
		return err
	}

	return nil
}

func (h *ExportFinalizeHandler) mergeShardCSVFiles(ctx context.Context, headers []string, shards []*model.ExportShardTask) (string, func(), error) {
	if h.storageSvc == nil || h.storageSvc.Provider() == nil {
		return "", nil, errors.New("对象存储服务未配置")
	}

	file, err := os.CreateTemp("", "export-merged-*.csv")
	if err != nil {
		return "", nil, err
	}

	cleanup := func() {
		_ = os.Remove(file.Name())
	}

	writer := csv.NewWriter(file)
	if err := writer.Write(headers); err != nil {
		file.Close()
		cleanup()
		return "", nil, err
	}

	for _, shard := range shards {
		if shard.Status != constants.ExportShardStatusSuccess {
			continue
		}
		if shard.FileKey == "" {
			file.Close()
			cleanup()
			return "", nil, fmt.Errorf("分片 %d 缺少文件Key", shard.ShardNo)
		}
		if err := h.appendShardFile(ctx, writer, shard.FileKey); err != nil {
			file.Close()
			cleanup()
			return "", nil, err
		}
	}

	writer.Flush()
	if err := writer.Error(); err != nil {
		file.Close()
		cleanup()
		return "", nil, err
	}
	if err := file.Sync(); err != nil {
		file.Close()
		cleanup()
		return "", nil, err
	}
	if err := file.Close(); err != nil {
		cleanup()
		return "", nil, err
	}

	return file.Name(), cleanup, nil
}

func (h *ExportFinalizeHandler) appendShardFile(ctx context.Context, writer *csv.Writer, key string) error {
	reader, err := h.storageSvc.Provider().Download(ctx, key)
	if err != nil {
		return err
	}
	defer reader.Close()

	csvReader := csv.NewReader(reader)
	for {
		row, err := csvReader.Read()
		if errors.Is(err, io.EOF) {
			break
		}
		if err != nil {
			return err
		}
		if err := writer.Write(row); err != nil {
			return err
		}
	}
	return nil
}

func prepareFinalExportFile(format, mergedCSVPath string) (string, int64, func(), error) {
	switch format {
	case constants.ExportTaskFormatCSV:
		stat, err := os.Stat(mergedCSVPath)
		if err != nil {
			return "", 0, nil, err
		}
		return mergedCSVPath, stat.Size(), func() {}, nil
	case constants.ExportTaskFormatXLSX:
		xlsxPath, fileSize, err := writeXLSXFromCSV(mergedCSVPath)
		if err != nil {
			return "", 0, nil, err
		}
		return xlsxPath, fileSize, func() { _ = os.Remove(xlsxPath) }, nil
	default:
		return "", 0, nil, fmt.Errorf("不支持的导出格式: %s", format)
	}
}

func (h *ExportFinalizeHandler) handleFinalizeError(ctx context.Context, taskID uint, updater uint, prefix string, err error) error {
	if isFinalRetry(ctx) {
		markErr := h.taskStore.MarkFailed(ctx, taskID, updater, prefix+": "+err.Error())
		if markErr != nil {
			return markErr
		}
		return nil
	}
	return err
}