import base64
import resumablehashlib
import json

from peewee import TextField


class ResumableSHAField(TextField):
  def db_value(self, value):
    sha_state = value.state()

    # One of the fields is a byte string, let's base64 encode it to make sure
    # we can store and fetch it regardless of default collocation.
    sha_state[3] = base64.b64encode(sha_state[3])

    return json.dumps(sha_state)

  def python_value(self, value):
    to_resume = resumablehashlib.sha256()
    if value is None:
      return to_resume

    sha_state = json.loads(value)

    # We need to base64 decode the data bytestring.
    sha_state[3] = base64.b64decode(sha_state[3])
    to_resume.set_state(sha_state)
    return to_resume


class JSONField(TextField):
  def db_value(self, value):
    return json.dumps(value)

  def python_value(self, value):
    if value is None or value == "":
      return {}
    return json.loads(value)