import { writeFileSync, mkdirSync, readFileSync, existsSync } from 'fs';
import { resolve, dirname, basename } from 'path';

const envPath = resolve(process.cwd(), '.env');
try {
  const envContent = readFileSync(envPath, 'utf-8');
  for (const line of envContent.split('\n')) {
    const match = line.match(/^([^#=]+)=(.*)$/);
    if (match && !process.env[match[1].trim()]) {
      process.env[match[1].trim()] = match[2].trim();
    }
  }
} catch {}


const API_BASE = 'https://api.banatie.app/api/v1';
const API_KEY = process.env.BANATIE_KEY || '';

const POLL_INTERVAL_MS = 2000;
const POLL_MAX_ATTEMPTS = 60; // 2 minutes total

async function pollGeneration(generationId) {
  for (let attempt = 1; attempt <= POLL_MAX_ATTEMPTS; attempt++) {
    const response = await fetch(`${API_BASE}/generations/${generationId}`, {
      headers: { 'X-API-Key': API_KEY },
    });

    if (!response.ok) {
      const text = await response.text();
      console.error(`Poll error ${response.status}: ${text}`);
      process.exit(1);
    }

    const result = await response.json();
    if (!result.success) {
      console.error(`Poll failed:`, result.error);
      process.exit(1);
    }

    const { status } = result.data;

    if (status === 'success') {
      return result.data;
    }

    if (status === 'failed') {
      console.error(`Generation failed: ${result.data.errorMessage || 'unknown error'}`);
      console.error('Suggestions: try a simpler prompt, check rate limits, or verify your BANATIE_KEY.');
      process.exit(1);
    }

    // Still pending/processing — wait and retry
    console.log(`Status: ${status} (attempt ${attempt}/${POLL_MAX_ATTEMPTS})...`);
    await new Promise(r => setTimeout(r, POLL_INTERVAL_MS));
  }

  console.error(`Generation timed out after ${POLL_MAX_ATTEMPTS * POLL_INTERVAL_MS / 1000}s. Try again or use a simpler prompt.`);
  process.exit(1);
}

async function uploadImage(filePath, { flowId, alias }) {
  const absolutePath = resolve(filePath);
  const fileBuffer = readFileSync(absolutePath);
  const fileName = basename(absolutePath);

  const ext = fileName.split('.').pop().toLowerCase();
  const mimeTypes = { png: 'image/png', jpg: 'image/jpeg', jpeg: 'image/jpeg', webp: 'image/webp' };
  const mime = mimeTypes[ext] || 'application/octet-stream';

  const form = new FormData();
  form.append('file', new Blob([fileBuffer], { type: mime }), fileName);
  form.append('alias', alias);
  if (flowId) form.append('flowId', flowId);

  console.log(`Uploading: ${filePath} as ${alias}${flowId ? ` (flow: ${flowId})` : ''}...`);

  const response = await fetch(`${API_BASE}/images/upload`, {
    method: 'POST',
    headers: { 'X-API-Key': API_KEY },
    body: form,
  });

  if (!response.ok) {
    const text = await response.text();
    console.error(`Upload error ${response.status}: ${text}`);
    if (response.status === 401) console.error('Check that BANATIE_KEY is set correctly in .env');
    if (response.status === 429) console.error('Rate limit exceeded. Wait before retrying (limit: 100 req/hour).');
    process.exit(1);
  }

  const result = await response.json();
  if (!result.success) {
    console.error(`Upload failed:`, result.error);
    process.exit(1);
  }

  console.log(`Uploaded: ${alias} (${result.data.id})`);
  return result.data;
}

function makeAlias(filePath, index) {
  const name = basename(filePath).replace(/\.[^.]+$/, '').replace(/[^a-z0-9-]/gi, '-').toLowerCase();
  return `@ref-${index + 1}-${name}`;
}

async function resolveRefs(refs) {
  if (!refs || refs.length === 0) return undefined;

  const aliases = [];
  let flowId = null;

  for (let i = 0; i < refs.length; i++) {
    const ref = refs[i];
    if (ref.startsWith('@')) {
      aliases.push(ref);
    } else if (existsSync(ref)) {
      const alias = makeAlias(ref, i);
      const data = await uploadImage(ref, { flowId, alias });
      if (!flowId) flowId = data.flowId;
      aliases.push(alias);
    } else {
      console.error(`Reference not found: ${ref}`);
      process.exit(1);
    }
  }

  return { referenceImages: aliases, flowId };
}

export async function generateImage({ prompt, output, aspectRatio = '1:1', refs, template, autoEnhance = true }) {
  if (!API_KEY) {
    console.error('BANATIE_KEY environment variable is not set.');
    console.error('Add BANATIE_KEY=your_key to the .env file in the project root.');
    process.exit(1);
  }

  const resolved = await resolveRefs(refs);

  const body = { prompt, aspectRatio, autoEnhance };

  if (template && autoEnhance) {
    body.enhancementOptions = { template };
  }

  if (resolved) {
    body.referenceImages = resolved.referenceImages;
    if (resolved.flowId) body.flowId = resolved.flowId;
  }

  const enhanceInfo = autoEnhance ? ` [template: ${template || 'general'}]` : ' [no enhance]';
  console.log(`Generating: "${prompt}" (${body.aspectRatio})${enhanceInfo}${resolved ? ` with ${resolved.referenceImages.length} ref(s)` : ''}...`);

  const response = await fetch(`${API_BASE}/generations`, {
    method: 'POST',
    headers: {
      'Content-Type': 'application/json',
      'X-API-Key': API_KEY,
    },
    body: JSON.stringify(body),
  });

  if (!response.ok) {
    const text = await response.text();
    console.error(`API error ${response.status}: ${text}`);
    if (response.status === 401) console.error('Check that BANATIE_KEY is set correctly in .env');
    if (response.status === 429) console.error('Rate limit exceeded. Wait before retrying (limit: 100 req/hour).');
    process.exit(1);
  }

  const result = await response.json();
  if (!result.success) {
    console.error(`Generation failed:`, result.error);
    process.exit(1);
  }

  // Handle async generation: poll if not yet complete
  let data = result.data;
  if (data.status === 'pending' || data.status === 'processing') {
    console.log(`Generation queued (id: ${data.id}), waiting for completion...`);
    data = await pollGeneration(data.id);
  }

  if (!data.outputImage?.storageUrl) {
    console.error('Generation completed but no output image found.');
    console.error('Response data:', JSON.stringify(data, null, 2));
    process.exit(1);
  }

  const imageUrl = data.outputImage.storageUrl;
  console.log(`Downloading from ${imageUrl}...`);

  const imageResponse = await fetch(imageUrl);
  if (!imageResponse.ok) {
    console.error(`Failed to download image: ${imageResponse.status}`);
    process.exit(1);
  }

  const buffer = Buffer.from(await imageResponse.arrayBuffer());
  const outputPath = resolve(output);
  mkdirSync(dirname(outputPath), { recursive: true });
  writeFileSync(outputPath, buffer);

  console.log(`Image saved: ${outputPath} (${data.outputImage.width}x${data.outputImage.height})`);
  return { path: outputPath, generation: data };
}

function parseArgs(args) {
  const result = { autoEnhance: true };
  for (let i = 0; i < args.length; i++) {
    if (args[i] === '--prompt') result.prompt = args[++i];
    else if (args[i] === '--output') result.output = args[++i];
    else if (args[i] === '--aspect-ratio') result.aspectRatio = args[++i];
    else if (args[i] === '--template') result.template = args[++i];
    else if (args[i] === '--no-enhance') result.autoEnhance = false;
    else if (args[i] === '--ref') {
      if (!result.refs) result.refs = [];
      result.refs.push(args[++i]);
    }
  }
  return result;
}

const args = parseArgs(process.argv.slice(2));
if (args.prompt && args.output) {
  generateImage(args);
} else if (process.argv.length > 2) {
  console.error('Usage: node banatie-gen.mjs --prompt "<description>" --output <path> [--aspect-ratio <ratio>] [--template <template>] [--no-enhance] [--ref <file|@alias>]...');
  process.exit(1);
}