import type {
  AssistantMessage,
  FunctionCall,
  JSONValue,
  StreamString
} from 'ai'

export interface StreamPart<CODE extends string, NAME extends string, TYPE> {
  code: CODE;
  name: NAME;
  parse: (value: JSONValue) => { type: NAME; value: TYPE };
}

const textStreamPart: StreamPart<'0', 'text', string> = {
  code: '0',
  name: 'text',
  parse: (value: JSONValue) => {
    if (typeof value !== 'string') {
      throw new Error('"text" parts expect a string value.')
    }
    return { type: 'text', value }
  }
}

const functionCallStreamPart: StreamPart<
  '1',
  'function_call',
  { function_call: FunctionCall }
> = {
  code: '1',
  name: 'function_call',
  parse: (value: JSONValue) => {
    if (
      value == null ||
      typeof value !== 'object' ||
      !('function_call' in value) ||
      typeof value.function_call !== 'object' ||
      value.function_call == null ||
      !('name' in value.function_call) ||
      !('arguments' in value.function_call) ||
      typeof value.function_call.name !== 'string' ||
      typeof value.function_call.arguments !== 'string'
    ) {
      throw new Error(
        '"function_call" parts expect an object with a "function_call" property.'
      )
    }

    return {
      type: 'function_call',
      value: value as unknown as { function_call: FunctionCall }
    }
  }
}

const dataStreamPart: StreamPart<'2', 'data', Array<JSONValue>> = {
  code: '2',
  name: 'data',
  parse: (value: JSONValue) => {
    if (!Array.isArray(value)) {
      throw new Error('"data" parts expect an array value.')
    }

    return { type: 'data', value }
  }
}

const errorStreamPart: StreamPart<'3', 'error', string> = {
  code: '3',
  name: 'error',
  parse: (value: JSONValue) => {
    if (typeof value !== 'string') {
      throw new Error('"error" parts expect a string value.')
    }
    return { type: 'error', value }
  }
}

const assistantMessage: StreamPart<'4', 'assistant_message', AssistantMessage> =
  {
    code: '4',
    name: 'assistant_message',
    parse: (value: JSONValue) => {
      if (
        value == null ||
        typeof value !== 'object' ||
        !('id' in value) ||
        !('role' in value) ||
        !('content' in value) ||
        typeof value.id !== 'string' ||
        typeof value.role !== 'string' ||
        value.role !== 'assistant' ||
        !Array.isArray(value.content) ||
        !value.content.every(
          item =>
            item != null &&
            typeof item === 'object' &&
            'type' in item &&
            item.type === 'text' &&
            'text' in item &&
            item.text != null &&
            typeof item.text === 'object' &&
            'value' in item.text &&
            typeof item.text.value === 'string'
        )
      ) {
        throw new Error(
          '"assistant_message" parts expect an object with an "id", "role", and "content" property.'
        )
      }

      return {
        type: 'assistant_message',
        value: value as AssistantMessage
      }
    }
  }

const assistantControlData: StreamPart<
  '5',
  'assistant_control_data',
  {
    threadId: string;
    messageId: string;
  }
> = {
  code: '5',
  name: 'assistant_control_data',
  parse: (value: JSONValue) => {
    if (
      value == null ||
      typeof value !== 'object' ||
      !('threadId' in value) ||
      !('messageId' in value) ||
      typeof value.threadId !== 'string' ||
      typeof value.messageId !== 'string'
    ) {
      throw new Error(
        '"assistant_control_data" parts expect an object with a "threadId" and "messageId" property.'
      )
    }

    return {
      type: 'assistant_control_data',
      value: {
        threadId: value.threadId,
        messageId: value.messageId
      }
    }
  }
}

const streamParts = [
  textStreamPart,
  functionCallStreamPart,
  dataStreamPart,
  errorStreamPart,
  assistantMessage,
  assistantControlData
] as const

// union type of all stream parts
type StreamParts =
  | typeof textStreamPart
  | typeof functionCallStreamPart
  | typeof dataStreamPart
  | typeof errorStreamPart
  | typeof assistantMessage
  | typeof assistantControlData;

/**
 * Maps the type of a stream part to its value type.
 */
type StreamPartValueType = {
  [P in StreamParts as P['name']]: ReturnType<P['parse']>['value'];
};

export type StreamPartType =
  | ReturnType<typeof textStreamPart.parse>
  | ReturnType<typeof functionCallStreamPart.parse>
  | ReturnType<typeof dataStreamPart.parse>
  | ReturnType<typeof errorStreamPart.parse>
  | ReturnType<typeof assistantMessage.parse>
  | ReturnType<typeof assistantControlData.parse>;

export const streamPartsByCode = {
  [textStreamPart.code]: textStreamPart,
  [functionCallStreamPart.code]: functionCallStreamPart,
  [dataStreamPart.code]: dataStreamPart,
  [errorStreamPart.code]: errorStreamPart,
  [assistantMessage.code]: assistantMessage,
  [assistantControlData.code]: assistantControlData
} as const

/**
 * The map of prefixes for data in the stream
 *
 * - 0: Text from the LLM response
 * - 1: (OpenAI) function_call responses
 * - 2: custom JSON added by the user using `Data`
 *
 * Example:
 * ```
 * 0:Vercel
 * 0:'s
 * 0: AI
 * 0: AI
 * 0: SDK
 * 0: is great
 * 0:!
 * 2: { "someJson": "value" }
 * 1: {"function_call": {"name": "get_current_weather", "arguments": "{\\n\\"location\\": \\"Charlottesville, Virginia\\",\\n\\"format\\": \\"celsius\\"\\n}"}}
 *```
 */
export const StreamStringPrefixes = {
  [textStreamPart.name]: textStreamPart.code,
  [functionCallStreamPart.name]: functionCallStreamPart.code,
  [dataStreamPart.name]: dataStreamPart.code,
  [errorStreamPart.name]: errorStreamPart.code,
  [assistantMessage.name]: assistantMessage.code,
  [assistantControlData.name]: assistantControlData.code
} as const

export const validCodes = streamParts.map(part => part.code)

/**
 * Parses a stream part from a string.
 *
 * @param line The string to parse.
 * @returns The parsed stream part.
 * @throws An error if the string cannot be parsed.
 */
export const parseStreamPart = (line: string): StreamPartType => {
  const firstSeparatorIndex = line.indexOf(':')

  if (firstSeparatorIndex === -1) {
    throw new Error('Failed to parse stream string. No separator found.')
  }

  const prefix = line.slice(0, firstSeparatorIndex)

  if (!validCodes.includes(prefix as keyof typeof streamPartsByCode)) {
    throw new Error(`Failed to parse stream string. Invalid code ${prefix}.`)
  }

  const code = prefix as keyof typeof streamPartsByCode

  const textValue = line.slice(firstSeparatorIndex + 1)
  const jsonValue: JSONValue = JSON.parse(textValue)

  return streamPartsByCode[code].parse(jsonValue)
}

/**
 * Prepends a string with a prefix from the `StreamChunkPrefixes`, JSON-ifies it,
 * and appends a new line.
 *
 * It ensures type-safety for the part type and value.
 */
export function formatStreamPart<T extends keyof StreamPartValueType> (
  type: T,
  value: StreamPartValueType[T]
): StreamString {
  const streamPart = streamParts.find(part => part.name === type)

  if (!streamPart) {
    throw new Error(`Invalid stream part type: ${type}`)
  }

  return `${streamPart.code}:${JSON.stringify(value)}\n`
}
