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

const envPath = resolve(dirname(new URL(import.meta.url).pathname), '../../.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) {
  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);

  console.log(`Uploading reference: ${filePath}...`);

  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: ${result.data.id}`);
  return result.data.id;
}

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

  const ids = [];
  for (const ref of refs) {
    if (existsSync(ref)) {
      ids.push(await uploadImage(ref));
    } else {
      ids.push(ref);
    }
  }
  return ids;
}

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 referenceImages = await resolveRefs(refs);
  const body = { prompt, aspectRatio };
  if (referenceImages) body.referenceImages = referenceImages;

  console.log(`Generating: "${prompt}" (${body.aspectRatio})${referenceImages ? ` with ${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.mjs --prompt "<description>" --output <path> [--aspect-ratio <ratio>] [--ref <file|id>]...');
  process.exit(1);
}
