import { parsePartialJson } from '@ai-sdk/ui-utils';
import { z } from 'zod';

export type DeepPartial<T> = T extends object
  ? T extends Array<any>
  ? T
  : { [P in keyof T]?: DeepPartial<T[P]> }
  : T;


export function removeJsonMarkdown(text: string) {
  if (text.startsWith('```json')) {
    text = text.slice(7);
  } else if (text.startsWith('json')) {
    text = text.slice(4);
  } else if (text.startsWith('```')) {
    text = text.slice(3);
  }
  if (text.endsWith('```')) {
    text = text.slice(0, -3);
  }
  return text;
}

/**
 * 解析流式的 JSON 数据
 * @param textStream 字符串流
 * @param schema zod schema 用于类型验证
 * @param isValid 自定义验证函数，用于判断解析出的 JSON 是否有效
 * @returns 异步生成器，yield 解析后的数据
 */
export async function* parseStreamingJson<T extends z.ZodType>(
  textStream: AsyncIterable<string>,
  schema: T,
  isValid: (value: DeepPartial<z.infer<T>>) => boolean
): AsyncGenerator<DeepPartial<z.infer<T>>> {
  let rawText = '';
  let isParseSuccessful = false;

  for await (const chunk of textStream) {
    rawText = removeJsonMarkdown(rawText + chunk);
    const parsed = parsePartialJson(rawText);

    isParseSuccessful = parsed.state === 'repaired-parse' || parsed.state === 'successful-parse';
    if (isParseSuccessful) {
      yield parsed.value as DeepPartial<z.infer<T>>;
    } else {
      console.dir(parsed, { depth: null, colors: true });
    }
  }

  return { isSuccessful: isParseSuccessful };
}