import { fetch } from 'undici';
import { resolve, dirname } from 'path';
import { writeFile, mkdir } from 'fs/promises';
import { createModuleLogger } from './utils/logger.js';
import { config } from './utils/config.js';
import { atomicWriteJSON, readJSON } from './modules/data-store/file-ops.js';
import { dataStore } from './modules/data-store/index.js';
import type { AppModules } from './wiring.js';

const log = createModuleLogger('kie-poller');

const KIE_API_BASE = 'https://api.kie.ai/api/v1';
const POLL_INTERVAL_MS = 10_000;

const FLUX_KONTEXT_MODELS = ['flux-kontext-pro', 'flux-kontext-max'];

export interface KIETrackedTask {
  kieTaskId: string;
  queueItemId: string;
  characterId: string;
  outputPath?: string;
  angleName?: string;
  isReferenceImage?: boolean;
  model: string;
  createdAt: string;
  status: 'pending' | 'completed' | 'failed';
}

interface TasksFile {
  tasks: KIETrackedTask[];
}

function getTasksFilePath(): string {
  return resolve(config.dataDir, 'queue', 'kie-tasks.json');
}

async function loadTasks(): Promise<KIETrackedTask[]> {
  const data = await readJSON<TasksFile>(getTasksFilePath());
  return data?.tasks ?? [];
}

async function saveTasks(tasks: KIETrackedTask[]): Promise<void> {
  await atomicWriteJSON(getTasksFilePath(), { tasks });
}

/**
 * Add a new KIE task to the tracking file for polling.
 */
export async function trackKIETask(task: KIETrackedTask): Promise<void> {
  const tasks = await loadTasks();
  tasks.push(task);
  await saveTasks(tasks);
  log.info({ kieTaskId: task.kieTaskId, model: task.model, characterId: task.characterId }, 'Tracking new KIE task');
}

function isFluxKontextModel(model: string): boolean {
  return FLUX_KONTEXT_MODELS.includes(model);
}

async function checkTaskStatus(task: KIETrackedTask): Promise<{ done: boolean; imageUrl?: string; failed?: boolean; failReason?: string }> {
  const url = isFluxKontextModel(task.model)
    ? `${KIE_API_BASE}/flux/kontext/record-info?taskId=${encodeURIComponent(task.kieTaskId)}`
    : `${KIE_API_BASE}/jobs/recordInfo?taskId=${encodeURIComponent(task.kieTaskId)}`;

  const response = await fetch(url, {
    method: 'GET',
    headers: {
      'Content-Type': 'application/json',
      Authorization: `Bearer ${config.kieApiKey}`,
    },
  });

  if (!response.ok) {
    log.warn({ kieTaskId: task.kieTaskId, status: response.status }, 'KIE status check failed');
    return { done: false };
  }

  const envelope = (await response.json()) as { code: number; data: any };
  if (envelope.code !== 200) {
    log.warn({ kieTaskId: task.kieTaskId, code: envelope.code }, 'KIE status returned non-200 code');
    return { done: false };
  }

  const data = envelope.data;

  if (isFluxKontextModel(task.model)) {
    // Flux kontext: successFlag 0=generating, 1=success, 2=create_failed, 3=generate_failed
    const flag = data.successFlag;
    if (flag === 1) {
      const imageUrl = data.response?.resultImageUrl;
      return { done: true, imageUrl };
    } else if (flag === 2 || flag === 3) {
      return { done: true, failed: true, failReason: data.errorMessage || `successFlag=${flag}` };
    }
    return { done: false };
  } else {
    // Market models: state = waiting/processing/completed/failed
    const state = data.state;
    log.debug({ kieTaskId: task.kieTaskId, state, model: task.model }, 'Market model status');
    if (state === 'success') {
      log.info({ kieTaskId: task.kieTaskId, resultJson: data.resultJson }, 'Task completed, raw resultJson');
      let imageUrl: string | undefined;
      if (data.resultJson) {
        try {
          const parsed = JSON.parse(data.resultJson);
          // KIE format: {"resultUrls":["https://..."]}
          if (typeof parsed === 'object' && !Array.isArray(parsed) && Array.isArray(parsed.resultUrls) && parsed.resultUrls.length > 0) {
            imageUrl = parsed.resultUrls[0];
          } else if (typeof parsed === 'string') {
            imageUrl = parsed;
          } else if (Array.isArray(parsed)) {
            imageUrl = parsed[0]?.url || parsed[0]?.image_url || parsed[0];
          } else if (typeof parsed === 'object') {
            imageUrl = parsed.url || parsed.image_url || parsed.resultImageUrl || parsed.output;
          }
        } catch {
          // resultJson might be a plain URL string (not JSON)
          imageUrl = data.resultJson;
        }
      }
      return { done: true, imageUrl };
    } else if (state === 'failed' || state === 'fail') {
      return { done: true, failed: true, failReason: data.failMsg || 'Task failed' };
    }
    return { done: false };
  }
}

async function downloadImage(imageUrl: string, destPath: string): Promise<void> {
  const dir = dirname(destPath);
  await mkdir(dir, { recursive: true });

  const response = await fetch(imageUrl);
  if (!response.ok) {
    throw new Error(`Failed to download image: ${response.status} ${response.statusText}`);
  }

  const buffer = Buffer.from(await response.arrayBuffer());
  await writeFile(destPath, buffer);
  log.info({ destPath, size: buffer.length }, 'Image downloaded successfully');
}

function resolveOutputPath(task: KIETrackedTask): string {
  if (task.outputPath) {
    return resolve(process.cwd(), task.outputPath);
  }
  return resolve(config.dataDir, 'characters', task.characterId, 'generated-images', `${task.kieTaskId}.jpg`);
}

async function pollOnce(): Promise<void> {
  const tasks = await loadTasks();
  const pending = tasks.filter(t => t.status === 'pending');

  if (pending.length === 0) return;

  log.info({ pendingCount: pending.length }, 'Polling KIE tasks');
  let changed = false;

  for (const task of pending) {
    try {
      const result = await checkTaskStatus(task);

      if (result.done) {
        changed = true;
        if (result.failed) {
          task.status = 'failed';
          log.warn({ kieTaskId: task.kieTaskId, reason: result.failReason }, 'KIE task failed');
        } else if (result.imageUrl) {
          const destPath = resolveOutputPath(task);
          await downloadImage(result.imageUrl, destPath);
          task.status = 'completed';
          log.info({ kieTaskId: task.kieTaskId, destPath, angleName: task.angleName }, 'KIE task completed and image saved');

          // Save reference image URL for character consistency
          if (task.isReferenceImage && task.outputPath) {
            const filename = task.outputPath.split('/').pop();
            if (filename) {
              await dataStore.addReferenceImageUrl(task.characterId, filename, result.imageUrl);
              log.info({ characterId: task.characterId, filename, url: result.imageUrl }, 'Reference image URL saved for consistency');
            }
          }
        } else {
          task.status = 'failed';
          log.warn({ kieTaskId: task.kieTaskId }, 'KIE task completed but no image URL found');
        }
      }
    } catch (err) {
      log.error({ kieTaskId: task.kieTaskId, err }, 'Error polling KIE task');
    }
  }

  if (changed) {
    await saveTasks(tasks);
  }
}

export function startKIEPoller(_modules: AppModules): { stop: () => void } {
  log.info({ intervalMs: POLL_INTERVAL_MS }, 'Starting KIE task poller');

  const interval = setInterval(async () => {
    try {
      await pollOnce();
    } catch (err) {
      log.error({ err }, 'KIE poller error');
    }
  }, POLL_INTERVAL_MS);

  // Run once immediately
  pollOnce().catch(err => log.error({ err }, 'KIE poller initial run error'));

  return {
    stop: () => {
      clearInterval(interval);
      log.info('KIE poller stopped');
    },
  };
}
