export type ChunkingStrategy = 'fixed' | 'overlap' | 'sentence' | 'paragraph' | 'markdown-heading' | 'semantic' | 'recursive';

export interface ChunkOpts {
  strategy: ChunkingStrategy;
  chunkSize?: number;   // token estimate target (1 token ≈ 4 chars). Default 2000.
  overlap?: number;     // overlap tokens for 'overlap' strategy. Default 200.
  separators?: string[]; // for 'recursive' strategy. Default ['\\n\\n', '\\n', '. ', ' ']
}

export interface Chunk {
  index: number;
  text: string;
  tokenEstimate: number;
  startChar: number;
  endChar: number;
  heading?: string;
}

export function estimateTokens(text: string): number {
  return Math.ceil(text.length / 4);
}

export function chunkContent(text: string, opts: ChunkOpts): Chunk[] {
  const size = opts.chunkSize ?? (parseInt(process.env.CRAWLER_CONTENT_CHUNK_SIZE ?? '') || 2000);
  const charSize = size * 4;
  const overlap = opts.overlap ?? (parseInt(process.env.CRAWLER_CONTENT_CHUNK_OVERLAP ?? '') || 200);
  const overlapChars = overlap * 4;

  switch (opts.strategy) {
    case 'fixed': return fixedChunk(text, charSize);
    case 'overlap': return overlapChunk(text, charSize, overlapChars);
    case 'sentence': return sentenceChunk(text, charSize);
    case 'paragraph': return paragraphChunk(text, charSize);
    case 'markdown-heading': return markdownHeadingChunk(text, charSize);
    case 'semantic': return semanticChunk(text, charSize);
    case 'recursive': return recursiveChunk(text, charSize, opts.separators ?? ['\n\n', '\n', '. ', ' ']);
    default: return fixedChunk(text, charSize);
  }
}

function makeChunk(text: string, index: number, startChar: number, heading?: string): Chunk {
  return { index, text, tokenEstimate: estimateTokens(text), startChar, endChar: startChar + text.length, heading };
}

function fixedChunk(text: string, charSize: number): Chunk[] {
  const chunks: Chunk[] = [];
  let i = 0;
  while (i < text.length) {
    const slice = text.slice(i, i + charSize);
    chunks.push(makeChunk(slice, chunks.length, i));
    i += charSize;
  }
  return chunks;
}

function overlapChunk(text: string, charSize: number, overlapChars: number): Chunk[] {
  const chunks: Chunk[] = [];
  let i = 0;
  while (i < text.length) {
    const slice = text.slice(i, i + charSize);
    chunks.push(makeChunk(slice, chunks.length, i));
    i += charSize - overlapChars;
    if (i <= 0) break;
  }
  return chunks;
}

function sentenceChunk(text: string, charSize: number): Chunk[] {
  const sentences = text.match(/[^.!?]*[.!?]+["']?\s*/g) ?? [text];
  return groupIntoChunks(sentences, charSize);
}

function paragraphChunk(text: string, charSize: number): Chunk[] {
  const paras = text.split(/\n\n+/).filter(p => p.trim());
  return groupIntoChunks(paras, charSize);
}

function markdownHeadingChunk(text: string, charSize: number): Chunk[] {
  const chunks: Chunk[] = [];
  const headingRe = /^(#{1,3}\s.+)$/m;
  const sections = text.split(/(?=^#{1,3}\s)/m).filter(s => s.trim());
  let pos = 0;
  for (const section of sections) {
    const headingMatch = section.match(headingRe);
    const heading = headingMatch ? headingMatch[1].replace(/^#+\s*/, '').trim() : undefined;
    if (section.length <= charSize) {
      chunks.push({ ...makeChunk(section.trim(), chunks.length, pos), heading });
    } else {
      // Section too large — sub-chunk by paragraph
      const sub = paragraphChunk(section, charSize);
      for (const s of sub) {
        chunks.push({ ...s, index: chunks.length, startChar: pos + s.startChar, endChar: pos + s.endChar, heading });
      }
    }
    pos += section.length;
  }
  return chunks.length > 0 ? chunks : fixedChunk(text, charSize);
}

function semanticChunk(text: string, charSize: number): Chunk[] {
  // Like paragraph but merges short paragraphs with the next
  const paras = text.split(/\n\n+/).filter(p => p.trim());
  const merged: string[] = [];
  let current = '';
  for (const para of paras) {
    if (current.length + para.length + 2 <= charSize) {
      current = current ? current + '\n\n' + para : para;
    } else {
      if (current) merged.push(current);
      current = para;
    }
  }
  if (current) merged.push(current);
  return groupIntoChunks(merged, charSize);
}

function recursiveChunk(text: string, charSize: number, separators: string[]): Chunk[] {
  if (text.length <= charSize) return [makeChunk(text, 0, 0)];
  const [sep, ...rest] = separators;
  if (!sep) return fixedChunk(text, charSize);
  const parts = text.split(sep).filter(p => p);
  const chunks: Chunk[] = [];
  let pos = 0;
  for (const part of parts) {
    if (part.length <= charSize) {
      chunks.push(makeChunk(part, chunks.length, pos));
    } else {
      const sub = recursiveChunk(part, charSize, rest);
      for (const s of sub) chunks.push({ ...s, index: chunks.length, startChar: pos + s.startChar, endChar: pos + s.endChar });
    }
    pos += part.length + sep.length;
  }
  return chunks;
}

function groupIntoChunks(parts: string[], charSize: number): Chunk[] {
  const chunks: Chunk[] = [];
  let current = '';
  let startChar = 0;
  let pos = 0;
  for (const part of parts) {
    if (current.length + part.length + 1 > charSize && current) {
      chunks.push(makeChunk(current.trim(), chunks.length, startChar));
      startChar = pos;
      current = part;
    } else {
      current = current ? current + ' ' + part : part;
    }
    pos += part.length + 1;
  }
  if (current.trim()) chunks.push(makeChunk(current.trim(), chunks.length, startChar));
  return chunks;
}

export function mergeSmallChunks(chunks: Chunk[], minTokens: number): Chunk[] {
  const minChars = minTokens * 4;
  const result: Chunk[] = [];
  let i = 0;
  while (i < chunks.length) {
    if (chunks[i].text.length < minChars && i + 1 < chunks.length) {
      const merged = chunks[i].text + '\n\n' + chunks[i + 1].text;
      result.push(makeChunk(merged, result.length, chunks[i].startChar, chunks[i].heading ?? chunks[i + 1].heading));
      i += 2;
    } else {
      result.push({ ...chunks[i], index: result.length });
      i++;
    }
  }
  return result;
}