// https://github.com/EvanZhouDev/bard-ai

class Bard {
  static JSON = 'json'
  static MD = 'markdown'

  // ID derived from Cookie
  SNlM0e

  // HTTPS Headers
  #headers

  // Resolution status of initialization call
  #initPromise

  #bardURL = 'https://bard.google.com'

  // Wether or not to log events to console
  #verbose = false

  // Fetch function
  #fetch = fetch

  constructor (cookie, config) {
    // Register some settings
    if (config?.verbose == true) this.#verbose = true
    if (config?.fetch) this.#fetch = config.fetch
    // 可变更访问地址，利用反向代理绕过区域限制
    if (config?.bardURL) this.#bardURL = config.bardURL

    // If a Cookie is provided, initialize
    if (cookie) {
      this.#initPromise = this.#init(cookie)
    } else {
      throw new Error('Please provide a Cookie when initializing Bard.')
    }
    this.cookie = cookie
  }

  // You can also choose to initialize manually
  async #init (cookie) {
    this.#verbose && console.log('🚀 Starting intialization')
    // Assign headers
    this.#headers = {
      Host: this.#bardURL.match(/^https?:\/\/([^\/]+)\/?$/)[1],
      'X-Same-Domain': '1',
      'User-Agent':
              'Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.114 Safari/537.36',
      'Content-Type': 'application/x-www-form-urlencoded;charset=UTF-8',
      Origin: this.#bardURL,
      Referer: this.#bardURL,
      Cookie: (typeof cookie === 'object') ? (Object.entries(cookie).map(([key, val]) => `${key}=${val};`).join('')) : ('__Secure-1PSID=' + cookie)
    }

    let responseText
    // Attempt to retrieve SNlM0e
    try {
      this.#verbose &&
              console.log('🔒 Authenticating your Google account')
      responseText = await this.#fetch(this.#bardURL, {
        method: 'GET',
        headers: this.#headers,
        credentials: 'include'
      })
        .then((response) => response.text())
    } catch (e) {
      // Failure to get server
      throw new Error(
        'Could not fetch Google Bard. You may be disconnected from internet: ' +
              e
      )
    }

    try {
      const SNlM0e = responseText.match(/SNlM0e":"(.*?)"/)[1]
      // Assign SNlM0e and return it
      this.SNlM0e = SNlM0e
      this.#verbose && console.log('✅ Initialization finished\n')
      return SNlM0e
    } catch {
      throw new Error(
        'Could not use your Cookie. Make sure that you copied correctly the Cookie with name __Secure-1PSID exactly. If you are sure your cookie is correct, you may also have reached your rate limit.'
      )
    }
  }

  async #uploadImage (name, buffer) {
    this.#verbose && console.log('🖼️ Starting image processing')
    let size = buffer.byteLength
    let formBody = [
          `${encodeURIComponent('File name')}=${encodeURIComponent([name])}`
    ]

    try {
      this.#verbose &&
              console.log('💻 Finding Google server destination')
      let response = await this.#fetch(
        'https://content-push.googleapis.com/upload/',
        {
          method: 'POST',
          headers: {
            'X-Goog-Upload-Command': 'start',
            'X-Goog-Upload-Protocol': 'resumable',
            'X-Goog-Upload-Header-Content-Length': size,
            'X-Tenant-Id': 'bard-storage',
            'Push-Id': 'feeds/mcudyrk2a4khkz'
          },
          body: formBody,
          credentials: 'include'
        }
      )

      const uploadUrl = response.headers.get('X-Goog-Upload-URL')
      this.#verbose && console.log('📤 Sending your image')
      response = await this.#fetch(uploadUrl, {
        method: 'POST',
        headers: {
          'X-Goog-Upload-Command': 'upload, finalize',
          'X-Goog-Upload-Offset': 0,
          'X-Tenant-Id': 'bard-storage'
        },
        body: buffer,
        credentials: 'include'
      })

      const imageFileLocation = await response.text()

      this.#verbose && console.log('✅ Image finished working\n')
      return imageFileLocation
    } catch (e) {
      throw new Error(
        'Could not fetch Google Bard. You may be disconnected from internet: ' +
              e
      )
    }
  }

  // Query Bard
  async #query (message, config) {
    let formatMarkdown = (text, images) => {
      if (!images) return text

      for (let imageData of images) {
        const formattedTag = `!${imageData.tag}(${imageData.url})`
        text = text.replace(
          new RegExp(`(?!\\!)\\[${imageData.tag.slice(1, -1)}\\]`),
          formattedTag
        )
      }

      return text
    }

    let { ids, imageBuffer } = config

    // Wait until after init
    await this.#initPromise

    this.#verbose && console.log('🔎 Starting Bard Query')

    // If user has not run init
    if (!this.SNlM0e) {
      throw new Error(
        "Please initialize Bard first. If you haven't passed in your Cookie into the class, run Bard.init(cookie)."
      )
    }

    this.#verbose && console.log('🏗️ Building Request')
    // HTTPS parameters
    const params = {
      bl: 'boq_assistant-bard-web-server_20230711.08_p0',
      _reqID: ids?._reqID ?? '0',
      rt: 'c'
    }

    // If IDs are provided, but doesn't have every one of the expected IDs, error
    const messageStruct = [
      [message],
      null,
      [null, null, null]
    ]

    if (imageBuffer) {
      let imageLocation = await this.#uploadImage(
        'bard-ai_upload',
        imageBuffer
      )
      messageStruct[0].push(0, null, [
        [[imageLocation, 1], 'bard-ai_upload']
      ])
    }

    if (ids) {
      const { conversationID, responseID, choiceID } = ids
      messageStruct[2] = [conversationID, responseID, choiceID]
    }

    // HTTPs data
    const data = {
      'f.req': JSON.stringify([null, JSON.stringify(messageStruct)]),
      at: this.SNlM0e
    }

    // URL that we are submitting to
    const url = new URL(
      '/_/BardChatUi/data/assistant.lamda.BardFrontendService/StreamGenerate',
      this.#bardURL
    )

    // Append parameters to the URL
    for (const key in params) {
      url.searchParams.append(key, params[key])
    }

    // Encode the data
    const formBody = Object.entries(data)
      .map(
        ([property, value]) =>
                  `${encodeURIComponent(property)}=${encodeURIComponent(
                      value
                  )}`
      )
      .join('&')

    this.#verbose && console.log('💭 Sending message to Bard')
    // Send the fetch request
    const chatData = await this.#fetch(url.toString(), {
      method: 'POST',
      headers: this.#headers,
      body: formBody,
      credentials: 'include'
    })
      .then((response) => {
        return response.text()
      })
      .then((text) => {
        return JSON.parse(text.split('\n')[3])[0][2]
      })
      .then((rawData) => JSON.parse(rawData))

    this.#verbose && console.log('🧩 Parsing output')
    // Get first Bard-recommended answer
    const answer = chatData[4][0]

    // Text of that answer
    const text = answer[1][0]

    // Get data about images in that answer
    const images =
          answer[4]?.map((x) => ({
            tag: x[2],
            url: x[3][0][0],
            info: {
              raw: x[0][0][0],
              source: x[1][0][0],
              alt: x[0][4],
              website: x[1][1],
              favicon: x[1][3]
            }
          })) ?? []

    this.#verbose && console.log('✅ All done!\n')
    // Put everything together and return
    return {
      content: formatMarkdown(text, images),
      images,
      ids: {
        conversationID: chatData[1][0],
        responseID: chatData[1][1],
        choiceID: answer[0],
        _reqID: String(parseInt(ids?._reqID ?? 0) + 100000)
      }
    }
  }

  async #parseConfig (config) {
    let result = {
      useJSON: false,
      imageBuffer: undefined, // Returns as {extension, filename}
      ids: undefined
    }

    // Verify that format is one of the two types
    if (config?.format) {
      switch (config.format) {
        case Bard.JSON:
          result.useJSON = true
          break
        case Bard.MD:
          result.useJSON = false
          break
        default:
          throw new Error(
            'Format can obly be Bard.JSON for JSON output or Bard.MD for Markdown output.'
          )
      }
    }

    // Verify that the image passed in is either a path to a jpeg, jpg, png, or webp, or that it is a Buffer
    if (config?.image) {
      if (
        config.image instanceof ArrayBuffer
      ) {
        result.imageBuffer = config.image
      } else if (
        typeof config.image === 'string' &&
              /\.(jpeg|jpg|png|webp)$/.test(config.image)
      ) {
        let fs

        try {
          fs = await import('fs')
        } catch {
          throw new Error(
            'Loading from an image file path is not supported in a browser environment.'
          )
        }

        result.imageBuffer = fs.readFileSync(
          config.image
        ).buffer
      } else {
        throw new Error(
          'Provide your image as a file path to a .jpeg, .jpg, .png, or .webp, or a Buffer.'
        )
      }
    }

    // Verify that all values in IDs exist
    if (config?.ids) {
      if (config.ids.conversationID && config.ids.responseID && config.ids.choiceID && config.ids._reqID) {
        result.ids = config.ids
      } else {
        throw new Error(
          'Please provide the IDs exported exactly as given.'
        )
      }
    }
    return result
  }

  // Ask Bard a question!
  async ask (message, config) {
    let { useJSON, imageBuffer, ids } = await this.#parseConfig(config)
    let response = await this.#query(message, { imageBuffer, ids })
    return useJSON ? response : response.content
  }

  createChat (ids) {
    let bard = this
    class Chat {
      ids = ids

      async ask (message, config) {
        let { useJSON, imageBuffer } = await bard.#parseConfig(config)
        let response = await bard.#query(message, {
          imageBuffer,
          ids: this.ids
        })
        this.ids = response.ids
        return useJSON ? response : response.content
      }

      export () {
        return this.ids
      }
    }

    return new Chat()
  }
}

export default Bard