// Images API - Extended OpenAI images with provider-specific parameters
// Endpoint: /images/generations

import type {
  ImageGenerateParams as OpenAIImageGenerateParams,
  ImagesResponse as OpenAIImagesResponse,
} from 'openai/resources/images';
import { type AIMLAPIClient } from '../client';
import { buildPath } from './_paths';

// ===== VERSION =====
const VERSION = 'v1' as const;

/**
 * Extended image generation parameters that support provider-specific options
 */
export interface ImageGenerateParams extends Omit<
  OpenAIImageGenerateParams,
  'size' | 'prompt' | 'output_format' | 'style'
> {
  /**
   * Provider-specific parameters
   */

  /**
   * Number of inference steps for models that support it (e.g., Flux, Stable Diffusion)
   */
  steps?: number;

  /**
   * Guidance scale for models that support it (e.g., Flux, Stable Diffusion)
   * Controls how much the image should adhere to the prompt
   */
  guidance_scale?: number;

  /**
   * Seed for reproducible generation
   */
  seed?: number;

  /**
   * Negative prompt for models that support it
   */
  negative_prompt?: string;

  /**
   * Width for models that support custom dimensions
   */
  width?: number;

  /**
   * Height for models that support custom dimensions
   */
  height?: number;

  /**
   * Aspect ratio for models that support it
   * General models: '1:1', '16:9', '9:16', etc.
   * KlingAI Image-O1: '21:9', '16:9', '4:3', '3:2', '1:1', '2:3', '3:4', '9:16'
   * Google Imagen 3 & 4: '1:1', '9:16', '16:9', '3:4', '4:3'
   * Google Gemini 2.5 Flash & 3 Pro: '21:9', '1:1', '4:3', '3:2', '2:3', '5:4', '4:5', '3:4', '16:9', '9:16'
   */
  aspect_ratio?:
    | '21:9'
    | '16:9'
    | '4:3'
    | '3:2'
    | '1:1'
    | '2:3'
    | '3:4'
    | '9:16'
    | '5:4'
    | '4:5'
    | string;

  /**
   * Enable safety checker for models that support it
   */
  safety_checker?: boolean;

  /**
   * LoRA identifier for models that support LoRA
   */
  lora?: string;

  /**
   * ControlNet parameters for models that support ControlNet
   */
  controlnet?: {
    type: string;
    image: string;
    strength?: number;
  };

  /**
   * The prompt for image generation (required for OpenAI compatibility)
   * For providers that don't require prompts (like Topaz Labs), use an empty string
   */
  prompt: string;

  /**
   * Image-to-image parameters for models that support img2img
   */
  image?: string; // Base64 encoded image or URL

  /**
   * Image URL for image editing models (reve/edit-image)
   * The URL of the reference image to be modified
   */
  image_url?: string;

  /**
   * The number of images to generate. Defaults to 1.
   */
  n?: number;

  /**
   * The format in which the generated images are returned.
   * Must be one of 'url' or 'b64_json'
   */
  response_format?: 'url' | 'b64_json';

  /**
   * Strength for image-to-image transformations
   */
  strength?: number;

  /**
   * Output format for image generation (provider-specific)
   * Options: 'jpeg', 'png', 'webp' (OpenAI compatible)
   * Note: This is for providers that use output_format instead of response_format
   */
  output_format?: 'jpeg' | 'png' | 'webp';

  /**
   * Quality setting for image generation models
   * OpenAI DALL-E: 'standard' | 'hd'
   * GPT Image-1: 'low' | 'medium' | 'high'
   */
  quality?: 'low' | 'medium' | 'high' | 'standard' | 'hd';

  /**
   * Image size for models that support custom dimensions or presets
   * For OpenAI compatibility, this matches OpenAI's exact size types
   * For provider-specific sizes, use image_size instead
   */
  size?:
    | 'auto'
    | '1024x1024'
    | '1536x1024'
    | '1024x1536'
    | '256x256'
    | '512x512'
    | '1792x1024'
    | '1024x1792'
    | null
    | undefined;

  /**
   * Provider-specific image size for custom dimensions or presets
   */
  image_size?:
    | {
        width?: number;
        height?: number;
      }
    | 'square_hd'
    | 'square'
    | 'portrait_4_3'
    | 'portrait_16_9'
    | 'landscape_4_3'
    | 'landscape_16_9'
    | '2K'
    | '4K'
    | string;

  /**
   * Enable prompt expansion (Hunyuan and Stable Diffusion v3 models)
   * If set to true, prompt will be upsampled with more details
   */
  prompt_expansion?: boolean;

  /**
   * Convert base64 to URL (reve/edit-image model)
   * If True, the URL to the image will be returned; otherwise, the file will be provided in base64 format
   */
  convert_base64_to_url?: boolean;

  /**
   * Background setting for GPT Image-1 model
   * Options: 'transparent', 'opaque', 'auto'
   */
  background?: 'transparent' | 'opaque' | 'auto';

  /**
   * Content moderation level for GPT Image-1 model
   * Options: 'low', 'auto'
   */
  moderation?: 'low' | 'auto';

  /**
   * Output compression for GPT Image-1 model (0-100)
   */
  output_compression?: number;

  /**
   * Resolution for KlingAI Image-O1 model
   * Options: '1K', '2K'
   */
  resolution?: '1K' | '2K';

  /**
   * Person generation setting for Google Imagen models
   * Options: 'dont_allow', 'allow_adult'
   */
  person_generation?: 'dont_allow' | 'allow_adult';

  /**
   * Safety setting for Google Imagen models
   * Options: 'block_low_and_above', 'block_medium_and_above', 'block_only_high'
   */
  safety_setting?: 'block_low_and_above' | 'block_medium_and_above' | 'block_only_high';

  /**
   * Add watermark setting for Google Imagen models
   */
  add_watermark?: boolean;

  /**
   * Enhanced resolution for Google Gemini models
   * Options: '1K', '2K', '4K'
   */
  enhanced_resolution?: '1K' | '2K' | '4K';

  /**
   * Safety tolerance level for Flux models
   * Options: '1', '2', '3', '4', '5', '6'
   * 1 being the most strict and 6 being the most permissive
   */
  safety_tolerance?: '1' | '2' | '3' | '4' | '5' | '6';

  /**
   * Style for image generation models
   * OpenAI DALL-E: 'vivid' | 'natural'
   * Recraft v3: 'any', 'realistic_image', 'digital_illustration', 'vector_illustration', and many sub-styles
   */
  style?:
    | 'vivid'
    | 'natural'
    | 'any'
    | 'realistic_image'
    | 'digital_illustration'
    | 'vector_illustration'
    | 'realistic_image/b_and_w'
    | 'realistic_image/hard_flash'
    | 'realistic_image/hdr'
    | 'realistic_image/natural_light'
    | 'realistic_image/studio_portrait'
    | 'realistic_image/enterprise'
    | 'realistic_image/motion_blur'
    | 'digital_illustration/pixel_art'
    | 'digital_illustration/hand_drawn'
    | 'digital_illustration/grain'
    | 'digital_illustration/infantile_sketch'
    | 'digital_illustration/2d_art_poster'
    | 'digital_illustration/handmade_3d'
    | 'digital_illustration/hand_drawn_outline'
    | 'digital_illustration/engraving_color'
    | 'digital_illustration/2d_art_poster_2'
    | 'vector_illustration/engraving'
    | 'vector_illustration/line_art'
    | 'vector_illustration/line_circuit'
    | 'vector_illustration/linocut';

  /**
   * Preferred colors for Recraft v3 model
   * Array of RGB color objects
   */
  colors?: Array<{
    r: number; // 0-255
    g: number; // 0-255
    b: number; // 0-255
  }>;

  /**
   * Number of images to generate (Alibaba models)
   */
  num_images?: number;

  /**
   * Enable safety checker (Alibaba models)
   */
  enable_safety_checker?: boolean;

  /**
   * Enable prompt expansion (Alibaba Z-Image models)
   */
  enable_prompt_expansion?: boolean;

  /**
   * Number of inference steps (Alibaba Z-Image models)
   */
  num_inference_steps?: number;

  /**
   * Acceleration mode (Alibaba Z-Image models)
   * Options: 'none', 'regular', 'high'
   */
  acceleration?: 'none' | 'regular' | 'high';

  /**
   * LoRA configurations (Alibaba Z-Image Turbo LoRA)
   */
  loras?: Array<{
    path: string;
    scale?: number;
  }>;

  /**
   * Image URLs for multi-image input (Alibaba Wan models)
   */
  image_urls?: string[];

  /**
   * Enhance prompt using LLM (Alibaba Wan models)
   */
  enhance_prompt?: boolean;

  /**
   * Add watermark (Alibaba models)
   */
  watermark?: boolean;

  /**
   * ByteDance-specific parameters
   */

  /**
   * Sync mode for ByteDance models (wait for generation)
   */
  sync_mode?: boolean;

  /**
   * Keep original size for ByteDance models
   */
  keep_size?: boolean;

  /**
   * Topaz Labs-specific parameters
   */

  /**
   * Sharpening mode for Topaz Labs Sharpen model
   * Options: 'Standard', 'Strong', 'Lens Blur', 'Lens Blur V2', 'Motion Blur', 'Natural', 'Refocus'
   */
  mode?:
    | 'Standard'
    | 'Strong'
    | 'Lens Blur'
    | 'Lens Blur V2'
    | 'Motion Blur'
    | 'Natural'
    | 'Refocus';

  /**
   * Subject detection for Topaz Labs models
   * Options: 'All', 'Foreground', 'Background'
   */
  subject_detection?: 'All' | 'Foreground' | 'Background';

  /**
   * Whether to enhance faces in the image
   */
  face_enhancement?: boolean;

  /**
   * Level of creativity for face enhancement (0-1)
   */
  face_enhancement_creativity?: number;

  /**
   * How sharp enhanced faces are relative to background (0-1)
   */
  face_enhancement_strength?: number;

  /**
   * Defines the overall intensity of the sharpening effect (0.01-1)
   */
  minor_denoise?: number;
}

/**
 * Image generation resource class that handles parameter transformation
 */
export class Images {
  constructor(private readonly _client: AIMLAPIClient) {}

  /**
   * Generate images using the extended parameter set
   * Transforms our extended interface to OpenAI's interface before making the request
   */
  async generate(params: ImageGenerateParams): Promise<OpenAIImagesResponse> {
    // Transform our extended interface to OpenAI's interface
    const openaiParams = this._transformToOpenAIParams(params);

    // Make the request using the transformed parameters
    return this._client.post<OpenAIImagesResponse>(buildPath('/images/generations', VERSION), {
      body: openaiParams,
    });
  }

  /**
   * Transform our extended ImageGenerateParams to OpenAI's ImageGenerateParams
   */
  private _transformToOpenAIParams(params: ImageGenerateParams): any {
    const transformed: any = {
      model: params.model,
      prompt: params.prompt,
    };

    // Add OpenAI-compatible parameters
    if (params.size !== undefined) {
      transformed.size = params.size;
    }

    if (params.output_format !== undefined) {
      transformed.response_format = params.output_format;
    }

    // Handle style conversion - only pass OpenAI-compatible styles
    if (params.style && this._isOpenAIStyle(params.style)) {
      transformed.style = params.style;
    }

    // Add other OpenAI parameters
    if (params.n !== undefined) transformed.n = params.n;
    if (params.quality !== undefined) transformed.quality = params.quality;
    if (params.response_format !== undefined) transformed.response_format = params.response_format;
    if (params.user !== undefined) transformed.user = params.user;

    // Add provider-specific parameters (they'll be passed through)
    const providerSpecificParams = {
      steps: params.steps,
      guidance_scale: params.guidance_scale,
      seed: params.seed,
      negative_prompt: params.negative_prompt,
      width: params.width,
      height: params.height,
      aspect_ratio: params.aspect_ratio,
      image_size: params.image_size,
      prompt_expansion: params.prompt_expansion,
      convert_base64_to_url: params.convert_base64_to_url,
      image_url: params.image_url,
      image: params.image, // Add missing image parameter
      style: params.style, // Include for provider-specific handling
      colors: params.colors,
      num_images: params.num_images,
      enable_safety_checker: params.enable_safety_checker,
      enable_prompt_expansion: params.enable_prompt_expansion,
      num_inference_steps: params.num_inference_steps,
      acceleration: params.acceleration,
      loras: params.loras,
      image_urls: params.image_urls,
      enhance_prompt: params.enhance_prompt,
      watermark: params.watermark,
      sync_mode: params.sync_mode,
      keep_size: params.keep_size,
      mode: params.mode,
      subject_detection: params.subject_detection,
      face_enhancement: params.face_enhancement,
      face_enhancement_creativity: params.face_enhancement_creativity,
      face_enhancement_strength: params.face_enhancement_strength,
      minor_denoise: params.minor_denoise,
    };

    // Handle image parameter transformation for specific models
    if (params.image && !params.image_url) {
      // For image-to-image models that expect image_url, transform image to image_url
      if (params.model && params.model.includes('/image-to-image')) {
        providerSpecificParams.image_url = params.image;
        delete providerSpecificParams.image;
      }
      // For Topaz Labs models that expect image_url, transform image to image_url
      else if (params.model && params.model.includes('topaz-labs/')) {
        providerSpecificParams.image_url = params.image;
        delete providerSpecificParams.image;
      }
      // For Alibaba image edit models that expect image_url, transform image to image_url
      else if (params.model && params.model.includes('alibaba/qwen-image-edit')) {
        providerSpecificParams.image_url = params.image;
        delete providerSpecificParams.image;
      }
    }

    // Filter out undefined provider-specific parameters
    Object.keys(providerSpecificParams).forEach((key) => {
      if (providerSpecificParams[key as keyof typeof providerSpecificParams] !== undefined) {
        transformed[key] = providerSpecificParams[key as keyof typeof providerSpecificParams];
      }
    });

    return transformed;
  }

  /**
   * Check if a style value is compatible with OpenAI's interface
   */
  private _isOpenAIStyle(style: string): style is 'vivid' | 'natural' {
    return style === 'vivid' || style === 'natural';
  }
}

/**
 * Image generation response type (re-export from OpenAI)
 */
export type { ImagesResponse as ImageResponse } from 'openai/resources/images';