import { isRagEnabled } from "./config";
import { knowledgeStoreExists } from "./knowledge-store";
import { formatKnowledgeForPrompt } from "./prompt";
import {
  buildGenerationRagQuery,
  hitsToCitations,
  resolvePrimaryRagPlatform,
  retrieveKnowledge,
} from "./retrieve";
import type { KnowledgePlatform, RagCitation } from "./types";

export interface GenerationRagContext {
  enabled: boolean;
  platform: KnowledgePlatform | null;
  captionBlock: string;
  imageBlock: string;
  videoBlock: string;
  citations: RagCitation[];
}

type RagInput = {
  platforms: string[];
  companyName: string;
  industry?: string;
  sequenceTitle?: string;
  problem?: string;
  impact?: string;
  contentCategory?: string;
  caption?: string;
};

function emptyContext(platform: KnowledgePlatform | null = null): GenerationRagContext {
  return {
    enabled: isRagEnabled(),
    platform,
    captionBlock: "",
    imageBlock: "",
    videoBlock: "",
    citations: [],
  };
}

/** Caption-phase RAG — before caption is generated. */
export async function retrieveCaptionKnowledge(
  input: RagInput,
): Promise<Pick<GenerationRagContext, "enabled" | "platform" | "captionBlock">> {
  if (!isRagEnabled()) return emptyContext();

  const platform = resolvePrimaryRagPlatform(input.platforms);
  const exists = await knowledgeStoreExists(platform);
  if (!exists) return { enabled: true, platform, captionBlock: "" };

  const baseQuery = buildGenerationRagQuery(input);
  const captionHits = await retrieveKnowledge({
    platform,
    query: baseQuery,
    topicBoost: ["content", "customer_audience"],
  });

  return {
    enabled: true,
    platform,
    captionBlock: formatKnowledgeForPrompt(captionHits, platform, "caption"),
  };
}

/** Media-phase RAG — after caption exists; aligns image/video KB with post content. */
export async function retrieveMediaKnowledge(
  input: RagInput & { caption: string },
): Promise<Pick<GenerationRagContext, "imageBlock" | "videoBlock" | "citations">> {
  if (!isRagEnabled()) {
    return { imageBlock: "", videoBlock: "", citations: [] };
  }

  const platform = resolvePrimaryRagPlatform(input.platforms);
  const exists = await knowledgeStoreExists(platform);
  if (!exists) return { imageBlock: "", videoBlock: "", citations: [] };

  const [imageHits, videoHits] = await Promise.all([
    retrieveKnowledge({
      platform,
      query: buildGenerationRagQuery({ ...input, mediaType: "image" }),
      topicBoost: "image",
    }),
    retrieveKnowledge({
      platform,
      query: buildGenerationRagQuery({ ...input, mediaType: "video" }),
      topicBoost: "video",
    }),
  ]);

  return {
    imageBlock: formatKnowledgeForPrompt(imageHits, platform, "image"),
    videoBlock: formatKnowledgeForPrompt(videoHits, platform, "video"),
    citations: hitsToCitations(
      [...imageHits, ...videoHits].filter(
        (h, i, arr) => arr.findIndex((x) => x.id === h.id) === i,
      ),
      platform,
    ),
  };
}

/** @deprecated Prefer retrieveCaptionKnowledge + retrieveMediaKnowledge */
export async function retrieveGenerationKnowledge(input: RagInput): Promise<GenerationRagContext> {
  const captionPhase = await retrieveCaptionKnowledge(input);
  const mediaPhase = input.caption
    ? await retrieveMediaKnowledge({ ...input, caption: input.caption })
    : { imageBlock: "", videoBlock: "", citations: [] as RagCitation[] };

  return {
    enabled: captionPhase.enabled,
    platform: captionPhase.platform,
    captionBlock: captionPhase.captionBlock,
    imageBlock: mediaPhase.imageBlock,
    videoBlock: mediaPhase.videoBlock,
    citations: mediaPhase.citations,
  };
}