/**
 * generationStore — SSE进度 / 图片候选 / 飞轮注入
 * image_candidate 无 batch 字段，所以 imageCandidates 平铺存储
 */
import { create } from 'zustand';
import { SseConnectionStatus } from '@/lib/sse';
import { TextCandidate, ImageCandidate } from '@/types/index';

interface FailedBatch {
  batch: string;  // 图片角色名（"hook"/"pain"等），对应 BE batch_failed.batch 字段
  reason: string;
  retryable: boolean;
}

interface GenerationState {
  // 连接状态
  connectionStatus: SseConnectionStatus;
  reconnectAttempt: number;

  // 文案生成进度
  textDone: number;
  textTotal: number;
  textCandidates: TextCandidate[];

  // 图片生成进度（平铺，不分 batch — 契约无 batch 字段）
  imageDone: number;
  imageTotal: number;
  imageCandidates: ImageCandidate[];
  failedBatches: FailedBatch[];

  // 飞轮注入（飞轮隐形，只显注入提示，无"训练AI"按钮）
  flywheelInjected: {
    recentPreference: string;
    rejectReasons: string[];
  } | null;

  setConnectionStatus: (s: SseConnectionStatus, attempt?: number) => void;
  setTextProgress: (done: number, total: number) => void;
  addTextCandidate: (candidate: TextCandidate) => void;
  patchTextCandidateScore: (candidateId: number, score: TextCandidate['score']) => void;
  setImageProgress: (done: number, total: number) => void;
  addImageCandidate: (candidate: ImageCandidate) => void;
  markBatchFailed: (batch: string, reason: string, retryable: boolean) => void;
  setFlywheelInjected: (data: { recentPreference: string; rejectReasons: string[] }) => void;
  hydrateFromTask: (data: {
    textCandidates: TextCandidate[];
    imageCandidates: ImageCandidate[];
  }) => void;
  reset: () => void;
}

export const useGenerationStore = create<GenerationState>((set) => ({
  connectionStatus: 'connecting',
  reconnectAttempt: 0,
  textDone: 0,
  textTotal: 0,
  textCandidates: [],
  imageDone: 0,
  imageTotal: 0,
  imageCandidates: [],
  failedBatches: [],
  flywheelInjected: null,

  setConnectionStatus(s, attempt = 0) {
    set({ connectionStatus: s, reconnectAttempt: attempt });
  },

  setTextProgress(done, total) {
    set({ textDone: done, textTotal: total });
  },

  addTextCandidate(candidate) {
    set((s) => ({ textCandidates: [...s.textCandidates, candidate] }));
  },

  patchTextCandidateScore(candidateId, score) {
    set((s) => ({
      textCandidates: s.textCandidates.map((tc) =>
        tc.candidate_id === candidateId ? { ...tc, score } : tc
      ),
    }));
  },

  setImageProgress(done, total) {
    set({ imageDone: done, imageTotal: total });
  },

  addImageCandidate(candidate) {
    set((s) => ({ imageCandidates: [...s.imageCandidates, candidate] }));
  },

  markBatchFailed(batch, reason, retryable) {
    set((s) => ({
      failedBatches: [...s.failedBatches, { batch, reason, retryable }],
    }));
  },

  setFlywheelInjected(data) {
    set({ flywheelInjected: data });
  },

  // 进入已完成/进行中任务时一次性回填（历史任务 SSE 不重放，必须主动拉接口）
  hydrateFromTask({ textCandidates, imageCandidates }) {
    set({
      textCandidates,
      textDone: textCandidates.length,
      textTotal: textCandidates.length || 0,
      imageCandidates,
      imageDone: imageCandidates.length,
      imageTotal: imageCandidates.length || 0,
    });
  },

  reset() {
    set({
      connectionStatus: 'connecting', reconnectAttempt: 0,
      textDone: 0, textTotal: 0, textCandidates: [],
      imageDone: 0, imageTotal: 0, imageCandidates: [],
      failedBatches: [], flywheelInjected: null,
    });
  },
}));
