import { fetch } from 'undici';
import { createModuleLogger } from '../../utils/logger.js';

const logger = createModuleLogger('gemini-client');

const GEMINI_API_BASE = 'https://geminiapi.codefast.app/v1';
const MAX_RETRIES = 3;
const INITIAL_BACKOFF_MS = 1000;
const POLL_INTERVAL_MS = 5000;
const MAX_POLL_ATTEMPTS = 60; // 5 minutes

export interface GeminiImageRequest {
  prompt: string;
  model?: 'gemini-3.0-pro' | 'gemini-3.1-flash';
  aspect_ratio?: 'landscape' | 'portrait' | 'square';
  seed?: number;
}

export interface GeminiImageVariationRequest {
  prompt: string;
  image: string; // base64
  model?: 'gemini-3.0-pro' | 'gemini-3.1-flash';
  mime_type?: string;
  aspect_ratio?: 'landscape' | 'portrait' | 'square';
  seed?: number;
}

interface GeminiJobResponse {
  projectId?: string;
  jobId: string;
  type: string;
  status: string;
  created_at?: string;
  queue?: {
    position?: number;
    ahead?: number;
    can_cancel?: boolean;
    status?: string;
    wait_seconds?: number;
  };
}

interface GeminiJobStatusResponse {
  jobId: string;
  job: {
    job_id: string;
    type: string;
    status: string;
    result?: {
      storage_urls?: string[];
      images?: string[];
    };
    error?: string | null;
    created_at?: string;
    updated_at?: string;
  };
  history_updated?: boolean;
  model?: string;
  queue?: {
    status?: string;
    position?: number;
    wait_seconds?: number;
  };
}

/**
 * Maps KIE aspect ratios (e.g. '9:16', '16:9', '1:1') to Gemini format.
 */
function mapAspectRatio(kieRatio?: string): 'landscape' | 'portrait' | 'square' {
  switch (kieRatio) {
    case '16:9':
    case '21:9':
    case '3:2':
      return 'landscape';
    case '9:16':
    case '3:4':
    case '2:3':
      return 'portrait';
    case '1:1':
    case '4:3':
      return 'square';
    default:
      return 'portrait'; // default for influencer photos
  }
}

export class GeminiClient {
  private readonly apiKey: string;
  private readonly initialBackoffMs: number;

  constructor(apiKey: string, options?: { initialBackoffMs?: number }) {
    this.apiKey = apiKey || '';
    this.initialBackoffMs = options?.initialBackoffMs ?? INITIAL_BACKOFF_MS;
  }

  /**
   * Creates a text-to-image job and polls until completion.
   * Returns the image URL.
   */
  async generateAndWait(request: GeminiImageRequest): Promise<string> {
    const job = await this.createImageJob(request);
    logger.info({ jobId: job.jobId, model: request.model }, 'Gemini image job created, polling...');
    return this.pollImageJob(job.jobId);
  }

  /**
   * Generates an image from prompt + optional reference images (KIE-compatible interface).
   * If reference image URLs are provided, uses the variation endpoint with the first reference.
   * Maps KIE aspect ratios to Gemini format.
   */
  async generateFromKIERequest(params: {
    prompt: string;
    aspectRatio?: string;
    model?: string;
    referenceImageUrls?: string[];
  }): Promise<string> {
    const aspect = mapAspectRatio(params.aspectRatio);

    // If reference images available, use variation endpoint for character consistency
    if (params.referenceImageUrls && params.referenceImageUrls.length > 0) {
      try {
        const refUrl = params.referenceImageUrls[0];
        logger.info({ refUrl: refUrl.slice(0, 80) }, 'Downloading reference image for Gemini variation');
        const base64 = await this.downloadImageAsBase64(refUrl);
        if (base64) {
          return this.generateVariationAndWait({
            prompt: params.prompt,
            image: base64.data,
            model: 'gemini-3.1-flash',
            mime_type: base64.mimeType,
            aspect_ratio: aspect,
          });
        }
      } catch (err: any) {
        logger.warn({ error: err.message }, 'Failed to use reference image, falling back to text-to-image');
      }
    }

    // Fallback: text-to-image without reference
    return this.generateAndWait({
      prompt: params.prompt,
      model: 'gemini-3.1-flash',
      aspect_ratio: aspect,
    });
  }

  /**
   * Creates an image variation job and polls until completion.
   */
  async generateVariationAndWait(request: GeminiImageVariationRequest): Promise<string> {
    const body: Record<string, unknown> = {
      prompt: request.prompt,
      image: request.image,
      model: request.model ?? 'gemini-3.1-flash',
      mime_type: request.mime_type ?? 'image/jpeg',
      aspect_ratio: request.aspect_ratio ?? 'portrait',
    };
    if (request.seed !== undefined) {
      body.seed = request.seed;
    }

    const job = await this.sendRequest<GeminiJobResponse>('POST', `${GEMINI_API_BASE}/image/variation`, body);
    logger.info({ jobId: job.jobId }, 'Gemini variation job created, polling...');
    return this.pollImageJob(job.jobId);
  }

  /**
   * Downloads an image from URL and returns base64 data + mime type.
   */
  private async downloadImageAsBase64(url: string): Promise<{ data: string; mimeType: string } | null> {
    try {
      const response = await fetch(url);
      if (!response.ok) return null;
      const buffer = Buffer.from(await response.arrayBuffer());
      const contentType = response.headers.get('content-type') || 'image/jpeg';
      const mimeType = contentType.split(';')[0].trim();
      return { data: buffer.toString('base64'), mimeType };
    } catch {
      return null;
    }
  }

  private async createImageJob(request: GeminiImageRequest): Promise<GeminiJobResponse> {
    const body: Record<string, unknown> = {
      prompt: request.prompt,
      model: request.model ?? 'gemini-3.1-flash',
      aspect_ratio: request.aspect_ratio ?? 'portrait',
    };
    if (request.seed !== undefined) {
      body.seed = request.seed;
    }

    return this.sendRequest<GeminiJobResponse>('POST', `${GEMINI_API_BASE}/image`, body);
  }

  private async pollImageJob(jobId: string): Promise<string> {
    for (let attempt = 1; attempt <= MAX_POLL_ATTEMPTS; attempt++) {
      await sleep(POLL_INTERVAL_MS);

      logger.debug({ jobId, attempt, maxAttempts: MAX_POLL_ATTEMPTS }, 'Polling Gemini image job');

      const status = await this.sendRequest<GeminiJobStatusResponse>(
        'POST',
        `${GEMINI_API_BASE}/image/status`,
        { job_id: jobId },
      );

      const jobStatus = status.job?.status;

      if (jobStatus === 'SUCCESS') {
        logger.info({ jobId, attempt, resultKeys: Object.keys(status.job?.result || {}) }, 'Gemini job SUCCESS, parsing result');

        const result = status.job?.result as Record<string, unknown> | undefined;
        if (result) {
          // Try storage_urls array
          const storageUrls = result.storage_urls as string[] | undefined;
          if (storageUrls && storageUrls.length > 0) {
            logger.info({ jobId, attempt, url: storageUrls[0] }, 'Gemini image job completed (storage_urls)');
            return storageUrls[0];
          }
          // Try storage_url (singular)
          if (typeof result.storage_url === 'string') {
            logger.info({ jobId, attempt, url: result.storage_url }, 'Gemini image job completed (storage_url)');
            return result.storage_url;
          }
          // Try url field directly
          if (typeof result.url === 'string') {
            logger.info({ jobId, attempt, url: result.url }, 'Gemini image job completed (url)');
            return result.url;
          }
          // Try images array (base64)
          const imgs = result.images as string[] | undefined;
          if (imgs && imgs.length > 0) {
            logger.info({ jobId, attempt }, 'Gemini image job completed (base64)');
            return `data:image/png;base64,${imgs[0]}`;
          }
          // Try image field (single base64)
          if (typeof result.image === 'string') {
            logger.info({ jobId, attempt }, 'Gemini image job completed (single base64)');
            return `data:image/png;base64,${result.image}`;
          }
        }

        // Last resort: log the full status for debugging
        logger.error({ jobId, fullResult: JSON.stringify(status).slice(0, 500) }, 'Gemini job SUCCESS but could not extract image URL');
        throw new Error(`Gemini job ${jobId} succeeded but no image URL found`);
      }

      if (jobStatus === 'ERROR' || jobStatus === 'CANCELED') {
        throw new Error(`Gemini job ${jobId} failed: ${status.job?.error || jobStatus}`);
      }

      // Still processing (QUEUED, PROCESSING, etc.)
      if (attempt % 12 === 0) {
        logger.info({ jobId, attempt, maxAttempts: MAX_POLL_ATTEMPTS, queuePosition: status.queue?.position },
          'Still waiting for Gemini image job...');
      }
    }

    throw new Error(`Gemini job ${jobId} timed out after ${MAX_POLL_ATTEMPTS * POLL_INTERVAL_MS / 1000} seconds`);
  }

  private async sendRequest<T>(method: 'GET' | 'POST', url: string, body?: unknown): Promise<T> {
    let lastError: Error | undefined;

    for (let attempt = 1; attempt <= MAX_RETRIES; attempt++) {
      try {
        logger.debug({ attempt, method, url }, 'Sending Gemini API request');

        const options: Record<string, unknown> = {
          method,
          headers: {
            'Content-Type': 'application/json',
            'Authorization': `Bearer ${this.apiKey}`,
          },
        };

        if (body !== undefined) {
          options.body = JSON.stringify(body);
        }

        const response = await fetch(url, options);

        if (!response.ok) {
          const errorText = await response.text();
          const err = new Error(
            `Gemini API error: ${response.status} ${response.statusText} - ${errorText}`,
          );
          logger.warn({ status: response.status, attempt, errorText }, 'Gemini API request failed');

          // Don't retry on 4xx (except 429)
          if (response.status >= 400 && response.status < 500 && response.status !== 429) {
            throw err;
          }
          lastError = err;
        } else {
          const data = (await response.json()) as T;
          logger.debug({ method, url }, 'Gemini API request successful');
          return data;
        }
      } catch (error) {
        if (lastError === undefined || lastError !== error) {
          lastError = error instanceof Error ? error : new Error(String(error));
        }
        if (lastError.message.includes('Gemini API error: 4') && !lastError.message.includes('429')) {
          throw lastError;
        }
        logger.warn({ attempt, error: lastError.message }, 'Gemini request attempt failed');
      }

      if (attempt < MAX_RETRIES) {
        const backoff = this.initialBackoffMs * Math.pow(2, attempt - 1);
        logger.debug({ backoff, attempt }, 'Waiting before retry');
        await sleep(backoff);
      }
    }

    logger.error({ error: lastError?.message }, 'All Gemini API retry attempts exhausted');
    throw lastError ?? new Error('Gemini API request failed after all retries');
  }
}

function sleep(ms: number): Promise<void> {
  return new Promise((resolve) => setTimeout(resolve, ms));
}

export function createGeminiClient(apiKey: string): GeminiClient {
  return new GeminiClient(apiKey);
}
