import { randomUUID } from "node:crypto";

import { decode, encode } from "./tokens";

interface StoredContinuation {
  remainingContent: string;
  metadata: Record<string, unknown>;
  timestamp: number;
}

export interface ChunkResult {
  content: string;
  chunk_number: number;
  has_more: boolean;
  continuation_token: string | null;
  tokens_in_chunk: number;
  remaining_tokens: number;
}

export class ChunkingManager {
  private readonly continuations = new Map<string, StoredContinuation>();

  constructor(private readonly ttl = 600) {}

  storeContinuation(remainingContent: string, metadata: Record<string, unknown>): string {
    const token = randomUUID();
    this.continuations.set(token, {
      remainingContent,
      metadata: { ...metadata },
      timestamp: Date.now() / 1000,
    });
    return token;
  }

  getNextChunk(token: string, chunkSize: number): ChunkResult | null {
    try {
      this.cleanupExpired();

      const entry = this.continuations.get(token);
      if (!entry) {
        return null;
      }

      if (Date.now() / 1000 - entry.timestamp > this.ttl) {
        this.continuations.delete(token);
        return null;
      }

      const tokens = encode(entry.remainingContent);
      const currentChunkNumber = Number((entry.metadata.chunk_number as number | undefined) ?? 1) + 1;

      if (tokens.length <= chunkSize) {
        const content = entry.remainingContent;
        this.continuations.delete(token);

        return {
          content,
          chunk_number: currentChunkNumber,
          has_more: false,
          continuation_token: null,
          tokens_in_chunk: encode(content).length,
          remaining_tokens: 0,
        };
      }

      const chunkTokens = tokens.slice(0, chunkSize);
      const remainingTokens = tokens.slice(chunkSize);
      const content = decode(chunkTokens);
      const newRemaining = decode(remainingTokens);

      const nextMetadata = {
        ...entry.metadata,
        chunk_number: currentChunkNumber,
      };

      const continuationToken = this.storeContinuation(newRemaining, nextMetadata);
      this.continuations.delete(token);

      return {
        content,
        chunk_number: currentChunkNumber,
        has_more: true,
        continuation_token: continuationToken,
        tokens_in_chunk: encode(content).length,
        remaining_tokens: remainingTokens.length,
      };
    } catch {
      return null;
    }
  }

  cleanupExpired(): number {
    const cutoff = Date.now() / 1000 - this.ttl;
    let removed = 0;

    for (const [token, entry] of this.continuations.entries()) {
      if (entry.timestamp < cutoff) {
        this.continuations.delete(token);
        removed += 1;
      }
    }

    return removed;
  }
}

export function getChunkSize(): number {
  const raw = process.env.RTFD_CHUNK_TOKENS ?? "2000";
  const parsed = Number.parseInt(raw, 10);
  return Number.isNaN(parsed) ? 2000 : parsed;
}