const { fromBuffer } = require('file-type');
const FormData = require("form-data");
const crypto = require("node:crypto");
const axios = require("axios");
const Jimp = require("jimp");

class GridPlus {
  constructor() {
    this.ins = axios.create({
      baseURL: 'https://api.grid.plus/v1',
      headers: {
        'user-agent': 'Mozilla/5.0 (Android 15; Mobile; SM-F958; rv:130.0) Gecko/130.0 Firefox/130.0',
        'X-AppID': '808645',
        'X-Platform': 'h5',
        'X-Version': '8.9.7',
        'X-SessionToken': '',
        'X-UniqueID': this.uid(),
        'X-GhostID': this.uid(),
        'X-DeviceID': this.uid(),
        'X-MCC': 'id-ID',
        sig: `XX${this.uid()+this.uid()}`
      }
    })
  }
  
  uid() {
    return crypto.randomUUID().replace(/-/g, '')
  }
  
  form(dt) {
    const form = new FormData();
    Object.entries(dt).forEach(([key, value]) => {
      form.append(key, String(value));
    });
    return form
  }
  
  async upload(buff, method) {
    try {
      if (!Buffer.isBuffer(buff)) throw 'data is not buffer!';
      const { mime, ext } = (await fromBuffer(buff)) || {};
      const d = await this.ins.post('/ai/web/nologin/getuploadurl', this.form({
        ext, method
      })).then(i => i.data);
      await axios.put(d.data.upload_url, buff, {
        headers: {
          'content-type': mime
        }
      });
      return d.data.img_url;
    } catch(e) {
      throw new Error('An error occurred while uploading the image');
    }
  }
  
  async task({ path, data, sl = () => false }) {
    const [start, interval, timeout] = [
      Date.now(), 3000, 60000
    ];
    return new Promise(async (resolve, reject) => {
      const check = async () => {
        if (Date.now() - start > timeout) {
          return reject(new Error(`Polling timed out for task`));
        };
        try {
          const dt = await this.ins({
            url: path,
            method: data ? 'POST' : 'GET',
            ...(data ? { data } : {})
          });
          if (!!dt.data.errmsg?.trim()) {
            reject(new Error(`Something a error with message: ${dt.data.errmsg}`));
          };
          if (!!sl(dt.data)) {
            return resolve(dt.data);
          };
          setTimeout(check, interval);
        } catch (error) {
          reject(error);
        }
      };
      check();
    });
  }
  
  async edit(buff, prompt) {
    try {
      const up = await this.upload(buff, 'wn_aistyle_nano');
      const dn = await this.ins.post('/ai/nano/upload', this.form({
        prompt, url: up
      })).then(i => i.data);
      if (!dn.task_id) throw 'taskId not found on request';
      const res = await this.task({
        path: `/ai/nano/get_result/${dn.task_id}`,
        sl: (dt) => dt.code === 0 && !!dt.image_url,
      });
      return res.image_url
    } catch(e) {
      throw new Error('Something error, message: ' + e.message);
    }
  }
}

async function createGrid(imageUrls) {
  const imageBuffers = await Promise.all(
    imageUrls.map(url => axios.get(url, { responseType: 'arraybuffer' }).then(r => Buffer.from(r.data)))
  );

  const images = await Promise.all(
    imageBuffers.map(async buffer => {
      const img = await Jimp.default ? Jimp.default.read(buffer) : new Promise((resolve, reject) => {
        new Jimp(buffer, (err, image) => {
          if (err) reject(err);
          else resolve(image);
        });
      });
      return img;
    })
  );

  const count = images.length;
  const targetSize = 1024;
  
  let cols, rows;
  if (count === 1) {
    const img = images[0];
    await img.resize(targetSize, targetSize);
    return await img.getBufferAsync(Jimp.MIME_JPEG);
  } else if (count === 2) {
    cols = 2;
    rows = 1;
  } else if (count === 3) {
    cols = 2;
    rows = 2;
  } else {
    cols = 2;
    rows = 2;
  }

  const cellSize = Math.floor(targetSize / cols);
  const gridHeight = cellSize * rows;
  const grid = await new Jimp(targetSize, gridHeight, 0xFFFFFFFF);

  for (let i = 0; i < count && i < 4; i++) {
    const col = i % cols;
    const row = Math.floor(i / cols);
    const x = col * cellSize;
    const y = row * cellSize;
    
    const resized = await images[i].cover(cellSize, cellSize);
    grid.composite(resized, x, y);
  }

  return await grid.getBufferAsync(Jimp.MIME_JPEG);
}

const handler = async (req, res) => {
  try {
    const { prompt, img_url1, img_url2, img_url3, img_url4, key } = req.query;

    if (!key) {
      return res.status(400).json({
        success: false,
        error: 'Missing required parameter: key'
      });
    }

    if (!prompt) {
      return res.status(400).json({
        success: false,
        error: 'Missing required parameter: prompt'
      });
    }

    const imageUrls = [img_url1, img_url2, img_url3, img_url4].filter(Boolean);

    if (imageUrls.length === 0) {
      return res.status(400).json({
        success: false,
        error: 'At least one image URL is required (img_url1)'
      });
    }

    if (imageUrls.length > 4) {
      return res.status(400).json({
        success: false,
        error: 'Maximum 4 images allowed'
      });
    }

    const imageBuffer = await createGrid(imageUrls);
    
    const gridPlus = new GridPlus();
    const result = await gridPlus.edit(imageBuffer, prompt);

    return res.json({
      author: "Herza",
      success: true,
      data: {
        prompt: prompt,
        image_url: result,
        original_urls: imageUrls,
        image_count: imageUrls.length
      },
      timestamp: new Date().toISOString()
    });

  } catch (error) {
    res.status(500).json({
      success: false,
      error: error.message,
      timestamp: new Date().toISOString()
    });
  }
};

module.exports = {
  name: 'Nano Banana',
  description: 'AI Image Editing using Grid Plus - Edit 1-4 images with text prompts',
  type: 'GET',
  routes: ['api/AI/nanobanana'],
  tags: ['ai', 'image', 'editing', 'nanobanana', 'multi-image'],
  parameters: ['prompt', 'img_url1', 'img_url2', 'img_url3', 'img_url4', 'key'],
  limit: 5,
  enabled: true,
  main: ['AI'],
  handler
};