const { CountTokensCommand } = require('@aws-sdk/client-bedrock-runtime');
const { BookStatsSchema } = require('./schemas');

/**
 * @param {string} text
 */
function estimateTokens(text) {
  return Math.ceil(String(text || '').length / 4);
}

/**
 * @param {{ send: Function }} client
 * @param {string} modelId
 * @param {string} text
 */
async function countWithBedrock(client, modelId, text) {
  const response = await client.send(
    new CountTokensCommand({
      modelId,
      input: {
        converse: {
          messages: [{ role: 'user', content: [{ text }] }],
        },
      },
    }),
  );
  if (typeof response.inputTokens !== 'number') {
    throw new Error('CountTokens did not return inputTokens');
  }
  return response.inputTokens;
}

/**
 * Count tokens for the marked book and each chapter.
 * Falls back to chars/4 when CountTokens is missing or fails.
 * @param {object} input
 */
async function countBook(input) {
  const { client, modelId, markedText, chapters, pages } = input;
  let tokenCountMethod = 'bedrock';
  let tokenCount;
  try {
    if (!client || !modelId) throw new Error('CountTokens client is not configured');
    tokenCount = await countWithBedrock(client, modelId, markedText);
  } catch {
    tokenCountMethod = 'estimate';
    tokenCount = estimateTokens(markedText);
  }

  const chapterCounts = [];
  for (const chapter of chapters) {
    const text = pages
      .filter((page) => page.page >= chapter.startPage && page.page <= chapter.endPage)
      .map((page) => page.text)
      .join('\n');
    if (tokenCountMethod === 'bedrock') {
      try {
        chapterCounts.push({
          chapterNo: chapter.chapterNo,
          tokenCount: await countWithBedrock(client, modelId, text),
          tokenCountMethod: 'bedrock',
        });
        continue;
      } catch {
        // This chapter falls back on its own.
      }
    }
    chapterCounts.push({
      chapterNo: chapter.chapterNo,
      tokenCount: estimateTokens(text),
      tokenCountMethod: 'estimate',
    });
  }

  return { tokenCount, tokenCountMethod, chapters: chapterCounts };
}

/**
 * @param {object} input
 */
function buildStats(input) {
  const { pages, markedText, tokenResult, cleaning, llm, timingsMs } = input;
  const stats = {
    pageCount: pages.length,
    charCount: markedText.length,
    wordCount: markedText.split(/\s+/).filter(Boolean).length,
    tokenCount: tokenResult.tokenCount,
    tokenCountMethod: tokenResult.tokenCountMethod,
    chapters: tokenResult.chapters,
    cleaning: {
      removedLineCount: cleaning.removedLineCount || 0,
      hyphenFixCount: cleaning.hyphenFixCount || 0,
    },
    llm: {
      inputTokens: llm.inputTokens || 0,
      outputTokens: llm.outputTokens || 0,
    },
    timingsMs,
  };
  return BookStatsSchema.parse(stats);
}

module.exports = { estimateTokens, countWithBedrock, countBook, buildStats };