import Replicate from "replicate";
import * as fs from "node:fs/promises";
import * as path from "node:path";

export interface ThumbnailGeneratorConfig {
  replicateApiKey: string;
}

export interface ThumbnailGenerateOptions {
  model?: string;
  systemPrompt?: string;
  imagePrompt?: string; // Path to local reference image
}

export class ThumbnailGenerator {
  private replicate: Replicate;

  constructor(config: ThumbnailGeneratorConfig) {
    this.replicate = new Replicate({
      auth: config.replicateApiKey,
    });
  }

  public async generateImage(
    topic: string,
    outputPath: string,
    options?: ThumbnailGenerateOptions,
  ): Promise<string> {
    const defaultSystemPrompt = `A highly polished, ultra-minimalist conceptual illustration. Style: high-end tech agency, clean modern 3D or flat vector art, extensive use of negative space, elegant monochrome palette (whites, light grays) with a single vibrant accent color (neon green or electric blue). Extremely clean and precise geometry. Absolutely no text, no photorealism, no chaotic lines, no messy sketches, no people.`;

    const systemPrompt = options?.systemPrompt || defaultSystemPrompt;
    const prompt = `${systemPrompt}\n\nTopic to illustrate abstractly: ${topic}`;

    console.log(`🎨 Generating thumbnail for topic: "${topic}"...`);

    const inputPayload: any = {
      prompt,
      aspect_ratio: "16:9",
      output_format: "png",
      output_quality: 90,
      prompt_upsampling: false,
    };

    if (options?.imagePrompt) {
      console.log(`🖼️ Using image style reference: ${options.imagePrompt}`);
      try {
        const absImgPath = path.isAbsolute(options.imagePrompt)
          ? options.imagePrompt
          : path.resolve(process.cwd(), options.imagePrompt);
        const imgBuffer = await fs.readFile(absImgPath);
        const base64 = imgBuffer.toString("base64");
        // Replicate models usually expect a data URI for image_prompt
        inputPayload.image_prompt = `data:image/png;base64,${base64}`;
      } catch (err) {
        console.warn(`⚠️ Could not load image prompt: ${err}`);
      }
    }

    // Default to the requested flux-1.1-pro model unless explicitly provided
    const model = options?.model || "black-forest-labs/flux-1.1-pro";

    const output = await this.replicate.run(model as `${string}/${string}`, {
      input: inputPayload,
    });

    // Replicate returns a ReadableStream for the output of flux-1.1-pro in newer Node SDKs
    // Or a string URL in older ones. We handle both.
    let buffer: Buffer;

    if (output instanceof ReadableStream) {
      console.log(`⬇️ Downloading generated stream from Replicate...`);
      const chunks: Uint8Array[] = [];
      const reader = output.getReader();
      while (true) {
        const { done, value } = await reader.read();
        if (done) break;
        if (value) chunks.push(value);
      }
      buffer = Buffer.concat(chunks);
    } else if (
      typeof output === "string" ||
      (Array.isArray(output) && typeof output[0] === "string")
    ) {
      const imageUrl = Array.isArray(output) ? output[0] : output;
      console.log(
        `⬇️ Downloading generated image from URL: ${imageUrl.substring(0, 50)}...`,
      );
      const response = await fetch(imageUrl);
      if (!response.ok) {
        throw new Error(`Failed to download image: ${response.statusText}`);
      }
      const arrayBuffer = await response.arrayBuffer();
      buffer = Buffer.from(arrayBuffer);
    } else if (Buffer.isBuffer(output)) {
      buffer = output;
    } else if (typeof output === "object") {
      console.log("Raw output object:", output);
      throw new Error("Unexpected output format from Replicate.");
    } else {
      throw new Error("Unknown output format from Replicate.");
    }

    const absPath = path.isAbsolute(outputPath)
      ? outputPath
      : path.resolve(process.cwd(), outputPath);
    await fs.mkdir(path.dirname(absPath), { recursive: true });
    await fs.writeFile(absPath, buffer);

    console.log(`✅ Saved thumbnail to: ${absPath}`);
    return absPath;
  }
}