export interface ThinkDelta {
  type: "reasoning" | "content";
  textDelta: string;
}

export class ThinkStreamParser {
  reasoning = "";
  content = "";

  private inThink = false;
  private buffer = "";

  private static readonly OPEN_TAG = "<think>";
  private static readonly CLOSE_TAG = "</think>";

  push(text: string): ThinkDelta[] {
    const deltas: ThinkDelta[] = [];
    this.buffer += text;

    while (this.buffer.length > 0) {
      if (this.inThink) {
        const closeIndex = this.buffer.indexOf(ThinkStreamParser.CLOSE_TAG);
        if (closeIndex !== -1) {
          const chunk = this.buffer.slice(0, closeIndex);
          if (chunk) {
            this.reasoning += chunk;
            deltas.push({ type: "reasoning", textDelta: chunk });
          }
          this.buffer = this.buffer.slice(
            closeIndex + ThinkStreamParser.CLOSE_TAG.length,
          );
          this.inThink = false;
          continue;
        }

        const safeLength = this.getSafeFlushLength(
          this.buffer,
          ThinkStreamParser.CLOSE_TAG,
        );
        if (safeLength > 0) {
          const chunk = this.buffer.slice(0, safeLength);
          this.reasoning += chunk;
          deltas.push({ type: "reasoning", textDelta: chunk });
          this.buffer = this.buffer.slice(safeLength);
        }
        break;
      }

      const openIndex = this.buffer.indexOf(ThinkStreamParser.OPEN_TAG);
      if (openIndex !== -1) {
        const chunk = this.buffer.slice(0, openIndex);
        if (chunk) {
          this.content += chunk;
          deltas.push({ type: "content", textDelta: chunk });
        }
        this.buffer = this.buffer.slice(
          openIndex + ThinkStreamParser.OPEN_TAG.length,
        );
        this.inThink = true;
        continue;
      }

      const safeLength = this.getSafeFlushLength(
        this.buffer,
        ThinkStreamParser.OPEN_TAG,
      );
      if (safeLength > 0) {
        const chunk = this.buffer.slice(0, safeLength);
        this.content += chunk;
        deltas.push({ type: "content", textDelta: chunk });
        this.buffer = this.buffer.slice(safeLength);
      }
      break;
    }

    return deltas;
  }

  flush(): ThinkDelta[] {
    if (!this.buffer) return [];

    const deltas: ThinkDelta[] = [];
    if (this.inThink) {
      this.reasoning += this.buffer;
      deltas.push({ type: "reasoning", textDelta: this.buffer });
    } else {
      this.content += this.buffer;
      deltas.push({ type: "content", textDelta: this.buffer });
    }
    this.buffer = "";
    return deltas;
  }

  private getSafeFlushLength(buffer: string, tag: string): number {
    for (
      let overlap = Math.min(buffer.length, tag.length - 1);
      overlap > 0;
      overlap--
    ) {
      if (buffer.endsWith(tag.slice(0, overlap))) {
        return buffer.length - overlap;
      }
    }
    return buffer.length;
  }
}