import pool from '../db';
import { extractKieTaskId, registerKieTask } from './kie-callbacks';

const KIE_API_KEY = process.env.KIE_API_KEY || '';
const KIE_MODEL = process.env.KIE_MODEL || 'grok-imagine/text-to-image';
const KIE_CREATE_URL = 'https://api.kie.ai/api/v1/jobs/createTask';
const KIE_STATUS_URL = 'https://api.kie.ai/api/v1/jobs/recordInfo';
const KIE_DOWNLOAD_URL = 'https://api.kie.ai/api/v1/common/download-url';
const KIE_POLL_INTERVAL = 2000;
const KIE_MAX_WAIT = 90000;
const KIE_CALLBACK_URL = process.env.KIE_CALLBACK_URL || 'https://rainmakersports.app/api/kie/callback';

interface GenerateKieImageOptions {
  aspectRatio?: string;
  category?: string;
  subcategory?: string;
  callbackUrl?: string;
}

export type KieCompletionSource = 'callback' | 'polling';

export interface KieImageResult {
  buffer: Buffer | null;
  completionSource: KieCompletionSource | null;
}

export async function generateKieImage(
  prompt: string,
  options: GenerateKieImageOptions = {}
): Promise<Buffer | null> {
  const result = await generateKieImageWithMeta(prompt, options);
  return result.buffer;
}

export async function generateKieImageWithMeta(
  prompt: string,
  options: GenerateKieImageOptions = {}
): Promise<KieImageResult> {
  if (!KIE_API_KEY) {
    return { buffer: null, completionSource: null };
  }

  const startTime = Date.now();
  const aspectRatio = options.aspectRatio || '16:9';
  const category = options.category || 'social_engine';
  const subcategory = options.subcategory || 'image_generation';
  const callbackUrl = options.callbackUrl || KIE_CALLBACK_URL;

  try {
    const createRes = await fetch(KIE_CREATE_URL, {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
        Authorization: `Bearer ${KIE_API_KEY}`,
      },
      body: JSON.stringify({
        model: KIE_MODEL,
        input: {
          prompt,
          aspect_ratio: aspectRatio,
        },
        ...(callbackUrl ? { callBackUrl: callbackUrl } : {}),
      }),
      signal: AbortSignal.timeout(15000),
    });

    if (!createRes.ok) {
      const errText = await createRes.text();
      throw new Error(`kie.ai create ${createRes.status}: ${errText.slice(0, 200)}`);
    }

    const createData: any = await createRes.json();
    if (createData.code !== 200) throw new Error(`kie.ai create error: ${createData.msg}`);

    const taskId = createData.data?.taskId;
    if (!taskId) throw new Error('No taskId returned from kie.ai');

    const completion = callbackUrl
      ? await waitForKieResult(taskId)
      : await pollForKieResult(taskId, 'polling');

    if (!completion?.resultUrl) throw new Error('kie.ai timed out waiting for image');

    let downloadUrl = completion.resultUrl;
    try {
      const dlRes = await fetch(KIE_DOWNLOAD_URL, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          Authorization: `Bearer ${KIE_API_KEY}`,
        },
        body: JSON.stringify({ url: completion.resultUrl }),
        signal: AbortSignal.timeout(10000),
      });

      if (dlRes.ok) {
        const dlData: any = await dlRes.json();
        if (dlData.code === 200 && dlData.data) {
          downloadUrl = dlData.data;
        }
      }
    } catch {
      // Use the original URL if download-url conversion fails.
    }

    const imgRes = await fetch(downloadUrl, {
      signal: AbortSignal.timeout(30000),
    });
    if (!imgRes.ok) throw new Error(`Image download failed: ${imgRes.status}`);

    const buffer = Buffer.from(await imgRes.arrayBuffer());

    pool.query(
      `INSERT INTO rm_api_usage (category, subcategory, provider, model, response_time_ms, success, metadata)
       VALUES ($1, $2, $3, $4, $5, true, $6)`,
      [
        category,
        subcategory,
        'kie.ai',
        KIE_MODEL,
        Date.now() - startTime,
        JSON.stringify({
          taskId,
          prompt_length: prompt.length,
          image_size: buffer.length,
          aspect_ratio: aspectRatio,
          callback_url: callbackUrl || null,
          completion_source: completion.completionSource,
        }),
      ]
    ).catch(() => {});

    return { buffer, completionSource: completion.completionSource };
  } catch (err: any) {
    pool.query(
      `INSERT INTO rm_api_usage (category, subcategory, provider, model, response_time_ms, success, error_message)
       VALUES ($1, $2, $3, $4, $5, false, $6)`,
      [category, subcategory, 'kie.ai', KIE_MODEL, Date.now() - startTime, err.message]
    ).catch(() => {});

    return { buffer: null, completionSource: null };
  }
}

async function waitForKieResult(taskId: string): Promise<{ resultUrl: string | null; completionSource: KieCompletionSource } | null> {
  try {
    const callbackPayload = await registerKieTask(taskId, KIE_MAX_WAIT);
    const callbackResult = extractResultUrlFromPayload(callbackPayload);
    if (callbackResult) {
      return { resultUrl: callbackResult, completionSource: 'callback' };
    }
  } catch {
    // Fall back to polling if callback never arrives or payload is incomplete.
  }

  return pollForKieResult(taskId, 'polling');
}

async function pollForKieResult(
  taskId: string,
  completionSource: KieCompletionSource
): Promise<{ resultUrl: string | null; completionSource: KieCompletionSource } | null> {
  const deadline = Date.now() + KIE_MAX_WAIT;

  while (Date.now() < deadline) {
    await new Promise(r => setTimeout(r, KIE_POLL_INTERVAL));

    const statusRes = await fetch(`${KIE_STATUS_URL}?taskId=${taskId}`, {
      headers: { Authorization: `Bearer ${KIE_API_KEY}` },
      signal: AbortSignal.timeout(10000),
    });

    if (!statusRes.ok) continue;

    const statusData: any = await statusRes.json();
    const task = statusData.data;

    if (task?.state === 'success') {
      return {
        resultUrl: extractResultUrlFromPayload(statusData),
        completionSource,
      };
    }

    if (task?.state === 'fail') {
      throw new Error(`kie.ai generation failed: ${task.failMsg || 'unknown'}`);
    }
  }

  return null;
}

function extractResultUrlFromPayload(payload: any): string | null {
  const task = payload?.data || payload;
  const resultJson = typeof task?.resultJson === 'string'
    ? JSON.parse(task.resultJson)
    : task?.resultJson;

  const urls = resultJson?.resultUrls || task?.resultUrls || payload?.resultUrls;
  if (Array.isArray(urls) && urls.length > 0) {
    return urls[0];
  }

  return payload?.resultUrl || task?.resultUrl || null;
}