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 || '';

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}`);
    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 }) {
  if (!API_KEY) {
    console.error('BANATIE_KEY environment variable is not set');
    process.exit(1);
  }

  const resolved = await resolveRefs(refs);
  const body = { prompt, aspectRatio };
  if (resolved) {
    body.referenceImages = resolved.referenceImages;
    if (resolved.flowId) body.flowId = resolved.flowId;
  }

  console.log(`Generating: "${prompt}" (${body.aspectRatio})${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}`);
    process.exit(1);
  }

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

  const imageUrl = result.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} (${result.data.outputImage.width}x${result.data.outputImage.height})`);
  return { path: outputPath, generation: result.data };
}

function parseArgs(args) {
  const result = {};
  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] === '--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>] [--ref <file|@alias>]...');
  process.exit(1);
}