package ai.decompile.ai.model.jdbc;

import java.sql.CallableStatement;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Types;
import java.util.Objects;
import org.hibernate.type.SqlTypes;
import org.hibernate.type.descriptor.ValueBinder;
import org.hibernate.type.descriptor.ValueExtractor;
import org.hibernate.type.descriptor.WrapperOptions;
import org.hibernate.type.descriptor.java.JavaType;
import org.hibernate.type.descriptor.jdbc.BasicBinder;
import org.hibernate.type.descriptor.jdbc.BasicExtractor;
import org.hibernate.type.descriptor.jdbc.JdbcType;
import org.postgresql.util.PGobject;

public class PostgreSQLVectorJdbcType implements JdbcType {
  @Override
  public int getJdbcTypeCode() {
    return Types.OTHER;
  }

  @Override
  public int getDefaultSqlTypeCode() {
    return SqlTypes.VECTOR;
  }

  @Override
  public String toString() {
    return "PostgreSQLVectorJdbcType";
  }

  @Override
  public <X> ValueBinder<X> getBinder(JavaType<X> javaType) {
    return new BasicBinder<>(javaType, this) {
      @Override
      protected void doBind(PreparedStatement st, X value, int index, WrapperOptions options)
          throws SQLException {
        var pgObj = new PGobject();
        pgObj.setType("vector");
        pgObj.setValue(formatVector((float[]) value));
        st.setObject(index, pgObj);
      }

      @Override
      protected void doBind(CallableStatement st, X value, String name, WrapperOptions options)
          throws SQLException {
        var pgObj = new PGobject();
        pgObj.setType("vector");
        pgObj.setValue(formatVector((float[]) value));
        st.setObject(name, pgObj);
      }
    };
  }

  @Override
  public <X> ValueExtractor<X> getExtractor(JavaType<X> javaType) {
    return new BasicExtractor<>(javaType, this) {
      @Override
      protected X doExtract(ResultSet rs, int paramIndex, WrapperOptions options)
          throws SQLException {
        return javaType.wrap(parseVector(rs.getString(paramIndex)), options);
      }

      @Override
      protected X doExtract(CallableStatement statement, int index, WrapperOptions options)
          throws SQLException {
        return javaType.wrap(parseVector(statement.getString(index)), options);
      }

      @Override
      protected X doExtract(CallableStatement statement, String name, WrapperOptions options)
          throws SQLException {
        return javaType.wrap(parseVector(statement.getString(name)), options);
      }
    };
  }

  static float[] parseVector(String vectorStr) {
    if (Objects.isNull(vectorStr) || vectorStr.isBlank()) {
      return new float[0];
    }

    var trimmed = vectorStr.trim();

    if (trimmed.startsWith("[") && trimmed.endsWith("]")) {
      trimmed = trimmed.substring(1, trimmed.length() - 1);
    }

    var parts = trimmed.split(",");
    var result = new float[parts.length];
    for (int i = 0; i < parts.length; i++) {
      result[i] = Float.parseFloat(parts[i].trim());
    }
    return result;
  }

  static String formatVector(float[] vector) {
    if (Objects.isNull(vector) || vector.length == 0) {
      return "[]";
    }

    var sb = new StringBuilder("[");
    for (int i = 0; i < vector.length; i++) {
      if (i > 0) {
        sb.append(",");
      }

      sb.append(vector[i]);
    }
    return sb.append("]").toString();
  }
}
