package org.apache.beam.sdk.extensions.sql.meta.provider.bigquery;

import static org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.rel2sql.SqlImplementor.POS;

import java.util.function.IntFunction;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.rel2sql.SqlImplementor;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rex.RexLiteral;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rex.RexNode;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rex.RexProgram;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlDialect;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlKind;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlLiteral;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlNode;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlWriter;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.parser.SqlParserPos;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.type.SqlTypeFamily;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.type.SqlTypeName;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.util.BitString;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.commons.lang.StringEscapeUtils;

public class BeamSqlUnparseContext extends SqlImplementor.SimpleContext {

  public BeamSqlUnparseContext(SqlDialect dialect, IntFunction<SqlNode> field) {
    super(dialect, field);
  }

  @Override
  public SqlNode toSql(RexProgram program, RexNode rex) {
    if (rex.getKind().equals(SqlKind.LITERAL)) {
      final RexLiteral literal = (RexLiteral) rex;
      if (literal.getTypeName().getFamily().equals(SqlTypeFamily.BINARY)) {
        BitString bitString = BitString.createFromBytes(literal.getValueAs(byte[].class));
        return new SqlByteStringLiteral(bitString, POS);
      } else if (literal.getTypeName().getFamily().equals(SqlTypeFamily.CHARACTER)) {
        String escaped = StringEscapeUtils.escapeJava(literal.getValueAs(String.class));
        return SqlLiteral.createCharString(escaped, POS);
      }
    }

    return super.toSql(program, rex);
  }

  private static class SqlByteStringLiteral extends SqlLiteral {

    SqlByteStringLiteral(BitString bytes, SqlParserPos pos) {
      super(bytes, SqlTypeName.BINARY, pos);
    }

    @Override
    public SqlByteStringLiteral clone(SqlParserPos pos) {
      return new SqlByteStringLiteral((BitString) this.value, pos);
    }

    @Override
    public void unparse(SqlWriter writer, int leftPrec, int rightPrec) {
      assert this.value instanceof BitString;

      StringBuilder builder = new StringBuilder("B'");
      for (byte b : ((BitString) this.value).getAsByteArray()) {
        builder.append(String.format("\\x%02X", b));
      }
      builder.append("'");

      writer.literal(builder.toString());
    }
  }
}

@@ -1737,20 +1737,39 @@
 ;%0A      
-if (
+SqlTypeFamily family = 
 literal.
@@ -1785,32 +1785,80 @@
 me().getFamily()
+;%0A      if (family != null) %7B%0A        if (family
 .equals(SqlTypeF
@@ -1870,24 +1870,26 @@
 .BINARY)) %7B%0A
+  
         BitS
@@ -1971,24 +1971,26 @@
 ));%0A        
+  
 return new S
@@ -2032,16 +2032,18 @@
 ;%0A      
+ 
+ 
 %7D else i
@@ -2049,41 +2049,14 @@
 if (
-literal.getTypeName().getF
+f
 amily
-()
 .equ
@@ -2091,24 +2091,26 @@
 ) %7B%0A        
+  
 String escap
@@ -2182,24 +2182,26 @@
 ));%0A        
+  
 return SqlLi
@@ -2234,24 +2234,34 @@
 aped, POS);%0A
+        %7D%0A
       %7D%0A    

Add null check for SqlTypeFamily