#!/usr/bin/env node
import puppeteer from 'puppeteer';
import axios from 'axios';
import sharp from 'sharp';
import PDFDocument from 'pdfkit';
import fs from 'fs';
import path from 'path';
import { fileURLToPath } from 'url';
import { load } from 'cheerio';
import { translate } from '@vitalets/google-translate-api';
import readline from 'readline';
import sizeOfImport from 'image-size';

// Resolve __dirname in ES modules
const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);

// Ensure image-size works regardless of export style
const sizeOf = sizeOfImport.default || sizeOfImport;

// ---- Constants ----
const WAIT_AFTER_NAVIGATION = 3000;  // ms to wait after navigation
const WAIT_AFTER_FLIP = 1000;        // ms to wait after each page flip
const MAX_ATTEMPTS = 50;             // maximum iterations if no new pages appear

// ---- Utility Functions ----
const sleep = (ms) => new Promise((resolve) => setTimeout(resolve, ms));

const promptInput = (query) => {
  const rl = readline.createInterface({
    input: process.stdin,
    output: process.stdout
  });
  return new Promise((resolve) => {
    rl.question(query, (answer) => {
      rl.close();
      resolve(answer.trim());
    });
  });
};

const cleanTitle = (title) =>
  title.replace(/[^\w\s-]/g, '').trim().replace(/\s+/g, '-');

const getFormattedDate = () => {
  const now = new Date();
  const pad = (num) => String(num).padStart(2, '0');
  return `${pad(now.getDate())}.${pad(now.getMonth() + 1)}.${now.getFullYear()}_${pad(now.getHours())}.${pad(now.getMinutes())}`;
};

// ---- Main Process ----
(async () => {
  // 1. Get URL from command line argument or prompt the user
  let url = process.argv[2];
  if (!url) {
    url = await promptInput('Please give me a URL for parsing from yunzhan: ');
    if (!url) {
      console.error('No URL provided. Exiting.');
      process.exit(1);
    }
  }

  // 2. Fix URL if needed
  if (url.includes('yunzhan365.com/basic/')) {
    console.log('Fixing the URL...');
    try {
      const resp = await axios.get(url);
      const $ = load(resp.data);
      const bookInfo = $('div.book-info');
      const titleElem = bookInfo.find('h1.title');
      const fixedUrl = titleElem.find('a').attr('href');
      if (fixedUrl) {
        url = fixedUrl;
        console.log('Fixed to ' + url);
      } else {
        console.error('Could not find the fixed URL.');
        process.exit(1);
      }
    } catch (err) {
      console.error('Error fetching the URL:', err);
      process.exit(1);
    }
  }

  // 3. Launch Puppeteer and navigate to the page
  console.log('Launching headless browser...');
  const browser = await puppeteer.launch({
    headless: true,
    args: ['--ignore-certificate-errors', '--disable-gpu', '--log-level=3']
  });
  const page = await browser.newPage();

  console.log('Navigating to page...');
  await page.goto(url, { waitUntil: 'networkidle2' });
  await sleep(WAIT_AFTER_NAVIGATION);

  // 4. Get total pages from the page's JS variable
  let numPages = await page.evaluate(() =>
    typeof originTotalPageCount === 'number' ? originTotalPageCount : null
  );
  if (!numPages) {
    console.log('Could not read originTotalPageCount; using fallback of 9999.');
    numPages = 9999;
  }
  console.log('Total pages (claimed):', numPages);

  // 5. Get and translate the page title
  let pageTitle = await page.title();
  console.log('Page title:', pageTitle);

  let translatedTitle = pageTitle;
  try {
    if (/[^\x00-\x7F]/.test(pageTitle)) {
      console.log('Translating title from Chinese (Simplified) to English...');
      const translation = await translate(pageTitle, { from: 'zh-CN', to: 'en' });
      translatedTitle = translation.text;
      console.log('Translated title:', translatedTitle);
    }
  } catch (err) {
    console.error('Translation error, using original title.', err);
  }

  console.log('Collecting images from DOM using page numbers...');

  // 6. Collect image URLs from the DOM using page numbers
  const pageMap = {};
  let attempts = 0;
  let prevCount = 0;

  while (attempts < MAX_ATTEMPTS) {
    const newPages = await page.evaluate(() => {
      const elems = Array.from(document.querySelectorAll('[id^="page"]'));
      return elems.map(el => {
        const num = parseInt(el.id.replace('page', ''), 10);
        const img = el.querySelector('img');
        const src = img ? img.src.split('?')[0] : null;
        return { num, src };
      }).filter(item => item.src !== null);
    });
  
    for (const p of newPages) {
      if (!pageMap[p.num]) {
        pageMap[p.num] = p.src;
      }
    }
    
    const currentCount = Object.keys(pageMap).length;
    console.log(`Total collected pages so far: ${currentCount}...`);
    
    if (currentCount > prevCount) {
      prevCount = currentCount;
      attempts = 0;
    } else {
      attempts++;
    }
    
    if (currentCount >= numPages) break;
    
    await page.evaluate(() => {
      if (typeof nextPageFun === 'function') {
        nextPageFun('mouse wheel flip');
      }
    });
    await sleep(WAIT_AFTER_FLIP);
  }

  let sortedPages = Object.keys(pageMap)
    .map(Number)
    .sort((a, b) => a - b);
  if (sortedPages.length > numPages) {
    sortedPages = sortedPages.slice(0, numPages);
  }
  const collectedImageURLs = sortedPages.map(num => pageMap[num]);
  console.log(`Collected ${collectedImageURLs.length} images in DOM order.`);

  await browser.close();

  // 7. Prepare PDF filename
  const safeTitle = cleanTitle(translatedTitle);
  const formattedDate = getFormattedDate();
  const resultsFolder = path.join(__dirname, 'results');
  if (!fs.existsSync(resultsFolder)) {
    fs.mkdirSync(resultsFolder);
  }
  const pdfPath = path.join(resultsFolder, `${safeTitle}_${formattedDate}.pdf`);

  // 8. Download images concurrently
  const downloadPromises = collectedImageURLs.map((link, idx) => {
    console.log(`Downloading page ${idx + 1} / ${collectedImageURLs.length}...`);
    return axios
      .get(link, { responseType: 'arraybuffer' })
      .then((res) => Buffer.from(res.data, 'binary'))
      .catch((err) => {
        console.error('Error downloading image:', link, err);
        return null;
      });
  });
  const downloaded = await Promise.all(downloadPromises);
  const imagesBuffers = downloaded.filter((buf) => buf !== null);
  if (imagesBuffers.length === 0) {
    console.error('No images were downloaded. Exiting.');
    process.exit(1);
  }

  // 9. Generate PDF using PDFKit
  console.log('Saving to PDF...');
  const doc = new PDFDocument({ autoFirstPage: false });
  const writeStream = fs.createWriteStream(pdfPath);
  doc.pipe(writeStream);

  for (const buffer of imagesBuffers) {
    let dimensions;
    try {
      dimensions = sizeOf(buffer);
    } catch (err) {
      console.error('Error getting image dimensions, skipping this page.', err);
      continue;
    }
    if (!['jpg', 'jpeg', 'png'].includes(dimensions.type)) {
      try {
        const converted = await sharp(buffer).png().toBuffer();
        dimensions = sizeOf(converted);
        doc.addPage({ size: [dimensions.width, dimensions.height] });
        doc.image(converted, 0, 0, {
          width: dimensions.width,
          height: dimensions.height
        });
      } catch (err) {
        console.error('Error converting image to PNG, skipping this page.', err);
      }
    } else {
      doc.addPage({ size: [dimensions.width, dimensions.height] });
      doc.image(buffer, 0, 0, {
        width: dimensions.width,
        height: dimensions.height
      });
    }
  }
  doc.end();

  await new Promise((resolve) => writeStream.on('finish', resolve));
  console.log('Done! PDF saved as', pdfPath);
})();

thanks, can you please support the url with a password to open? such as https://book.yunzhan365.com/irtl/ycmy/mobile/index.html,the pwd is 999999