import type {
  PaperChunk,
  ScientificBlock,
  ScientificBlockKind,
} from './types';

const TARGET_TOKENS = 650;
const MIN_TOKENS = 400;
const MAX_TOKENS = 900;
const OVERLAP_TOKENS = 70;

export function estimateScientificTokens(value: string): number {
  return value.match(/[\\\p{L}\p{N}_^+.-]+|[^\s]/gu)?.length ?? 0;
}

function hashString(value: string): string {
  // Stable deterministic ID helper. Cryptographic content hashes are filled
  // by sha256Text below; this is only used while composing chunk identity.
  let hash = 2166136261;
  for (let index = 0; index < value.length; index += 1) {
    hash ^= value.charCodeAt(index);
    hash = Math.imul(hash, 16777619);
  }
  return (hash >>> 0).toString(16).padStart(8, '0');
}

export async function sha256Text(value: string): Promise<string> {
  const bytes = new TextEncoder().encode(value);
  const digest = await crypto.subtle.digest('SHA-256', bytes);
  return [...new Uint8Array(digest)]
    .map((byte) => byte.toString(16).padStart(2, '0'))
    .join('');
}

function sentenceParts(text: string): string[] {
  const parts = text.match(/[^.!?\n]+(?:[.!?]+(?=\s|$)|$)/gu)?.map((part) => part.trim()).filter(Boolean);
  return parts?.length ? parts : [text.trim()];
}

function splitOversizedBlock(block: ScientificBlock): ScientificBlock[] {
  if (estimateScientificTokens(block.text) <= MAX_TOKENS) return [block];
  const result: ScientificBlock[] = [];
  let current: string[] = [];
  let tokens = 0;
  for (const sentence of sentenceParts(block.text)) {
    const next = estimateScientificTokens(sentence);
    if (current.length && tokens + next > MAX_TOKENS) {
      const text = current.join(' ').trim();
      result.push({ ...block, id: `${block.id}:${result.length}`, text });
      current = [];
      tokens = 0;
    }
    current.push(sentence);
    tokens += next;
  }
  if (current.length) result.push({ ...block, id: `${block.id}:${result.length}`, text: current.join(' ').trim() });
  return result;
}

function chunkKind(blocks: readonly ScientificBlock[]): Exclude<ScientificBlockKind, 'heading'> {
  const preferred = blocks.find((block) => block.kind !== 'paragraph' && block.kind !== 'heading');
  return (preferred?.kind ?? 'paragraph') as Exclude<ScientificBlockKind, 'heading'>;
}

function trimOverlap(blocks: readonly ScientificBlock[]): ScientificBlock[] {
  const result: ScientificBlock[] = [];
  let total = 0;
  for (let index = blocks.length - 1; index >= 0; index -= 1) {
    const block = blocks[index]!;
    const tokens = estimateScientificTokens(block.text);
    if (result.length && total + tokens > OVERLAP_TOKENS) break;
    result.unshift(block);
    total += tokens;
    if (total >= OVERLAP_TOKENS) break;
  }
  return result;
}

export async function chunkScientificBlocks(
  paperId: string,
  sourceBlocks: readonly ScientificBlock[],
): Promise<PaperChunk[]> {
  const blocks = sourceBlocks
    .filter((block) => block.kind !== 'heading' && block.text.trim())
    .flatMap(splitOversizedBlock);
  const chunks: PaperChunk[] = [];
  let current: ScientificBlock[] = [];
  let tokens = 0;

  const flush = async () => {
    if (!current.length) return;
    const text = current.map((block) => block.text.trim()).filter(Boolean).join('\n\n');
    const pageStart = Math.min(...current.map((block) => block.pageStart));
    const pageEnd = Math.max(...current.map((block) => block.pageEnd));
    const section = [...current].reverse().find((block) => block.section)?.section;
    const heading = [...current].reverse().find((block) => block.heading)?.heading;
    const contentHash = await sha256Text(`${chunkKind(current)}\n${section ?? ''}\n${text}`);
    chunks.push({
      id: `${paperId}:chunk:${chunks.length}:${hashString(contentHash)}`,
      paperId,
      kind: chunkKind(current),
      text,
      pageStart,
      pageEnd,
      ...(section ? { section } : {}),
      ...(heading ? { heading } : {}),
      contentHash,
    });
  };

  for (const block of blocks) {
    const blockTokens = estimateScientificTokens(block.text);
    const sectionChanged = current.length > 0 && block.section && current.at(-1)?.section && block.section !== current.at(-1)?.section;
    if (current.length && (tokens + blockTokens > MAX_TOKENS || (sectionChanged && tokens >= MIN_TOKENS))) {
      await flush();
      current = trimOverlap(current);
      tokens = current.reduce((sum, item) => sum + estimateScientificTokens(item.text), 0);
    }
    current.push(block);
    tokens += blockTokens;
    if (tokens >= TARGET_TOKENS) {
      await flush();
      current = trimOverlap(current);
      tokens = current.reduce((sum, item) => sum + estimateScientificTokens(item.text), 0);
    }
  }
  if (current.length) await flush();
  return chunks;
}