import { useEffect, useMemo, useRef, useState } from 'react'
import { MarkdownBody } from './markdown'

function segmentForStreaming(text: string): string[] {
  const tokens: string[] = []
  let i = 0
  const isCJK = (c: string) => /[\u4e00-\u9fff\u3040-\u30ff]/.test(c)
  const isWordChar = (c: string) => /[a-zA-Z0-9']/.test(c)
  while (i < text.length) {
    const c = text[i]!
    if (isCJK(c)) {
      tokens.push(c)
      i += 1
    } else if (isWordChar(c)) {
      let end = i
      while (end < text.length && isWordChar(text[end]!)) end += 1
      tokens.push(text.slice(i, end))
      i = end
    } else {
      let end = i
      while (end < text.length && !isCJK(text[end]!) && !isWordChar(text[end]!)) end += 1
      tokens.push(text.slice(i, end))
      i = end
    }
  }
  return tokens
}

function splitCompletedBlocks(visibleText: string): { completed: string; partial: string } {
  const blocks = visibleText.split(/\n\n+/)
  if (blocks.length <= 1) return { completed: '', partial: visibleText }
  const partial = blocks.pop() ?? ''
  const completed = blocks.join('\n\n')
  return { completed, partial }
}

export interface StreamingMarkdownBlockProps {
  content: string
  isActive: boolean
  onStreamComplete?: () => void
  tokenDelayMs?: number
  className?: string
}

const TOKENS_PER_TICK = 3

export function StreamingMarkdownBlock({
  content,
  isActive,
  onStreamComplete,
  tokenDelayMs = 14,
  className,
}: StreamingMarkdownBlockProps) {
  const tokens = useMemo(() => segmentForStreaming(content), [content])
  const [streamState, setStreamState] = useState({ content, visibleCount: isActive ? 0 : tokens.length })
  const onStreamCompleteRef = useRef(onStreamComplete)

  useEffect(() => {
    onStreamCompleteRef.current = onStreamComplete
  }, [onStreamComplete])

  const visibleCount = streamState.content === content
    ? streamState.visibleCount
    : isActive
      ? 0
      : tokens.length
  const visibleText = useMemo(() => tokens.slice(0, visibleCount).join(''), [tokens, visibleCount])
  const { completed, partial } = useMemo(() => splitCompletedBlocks(visibleText), [visibleText])

  useEffect(() => {
    if (!isActive) return
    const len = tokens.length
    if (len === 0) {
      onStreamCompleteRef.current?.()
      return
    }
    let k = 0
    const id = setInterval(() => {
      k = Math.min(k + TOKENS_PER_TICK, len)
      setStreamState({ content, visibleCount: k })
      if (k >= len) {
        clearInterval(id)
        onStreamCompleteRef.current?.()
      }
    }, tokenDelayMs)
    return () => clearInterval(id)
  }, [isActive, content, tokens.length, tokenDelayMs])

  const isStreaming = isActive && visibleCount < tokens.length

  if (tokens.length === 0) return null

  if (isStreaming) {
    return (
      <div className={className}>
        {completed ? (
          <div className="streaming-completed">
            <MarkdownBody>{completed}</MarkdownBody>
          </div>
        ) : null}
        {partial ? (
          <span className="streaming-partial text-sm leading-sm text-text streaming-token-fade-in">
            {partial}
          </span>
        ) : null}
      </div>
    )
  }

  return (
    <div className={className}>
      <MarkdownBody>{content}</MarkdownBody>
    </div>
  )
}