/*
 * RiskScape™ Copyright New Zealand Institute for Earth Science Limited
 * (Earth Sciences New Zealand) is distributed for research purposes only
 * under the terms of AGPLv3.
 *
 * RiskScape™ Copyright 2026 New Zealand Institute for Earth Science
 * Limited (Earth Sciences New Zealand). All rights reserved. Source code
 * available under the AGPLv3.
 * 
 * This program is free software: you can redistribute it and/or modify it under
 *  the terms of the GNU Affero General Public License as published by the Free
 *  Software Foundation, either version 3 of the License, or (at your option) any
 *  later version.
 * 
 * This program is distributed for RESEARCH PURPOSES ONLY, in the hope that it will
 * be useful for research and education initiatives.
 * 
 * If you are not a researcher, or you are a researcher who wishes to use this
 * program on terms other than AGPLv3 (including those who wish to restrict the
 * distribution of any source code created using this program), please contact:
 * https://riskscape.org.nz
 * 
 * This program is distributed WITHOUT ANY WARRANTY; without even the implied
 * warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Affero General Public License for more details.  You should have received a copy
 * of the GNU Affero General Public License along with this program.  If not, see
 * <http://www.gnu.org/licenses/>.
 * 
 * By way of summary only, under the AGPLv3:
 *     • Permissions of this strongest copyleft license are conditioned
 *       on making available complete source code of licensed works and
 *       modifications, which include larger works using a licensed work,
 *       under the same license.
 *     • Copyright and license notices must be preserved.
 *     • Contributors provide an express grant of patent rights.
 *     • When a modified version is used to provide a service over a
 *       network, the complete source code of the modified version must be made
 *       available.
 */
package nz.org.riskscape.engine.defaults.data.jdbc;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.UnsupportedEncodingException;
import java.net.URI;
import java.net.URLEncoder;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Function;
import java.util.stream.Collectors;

import javax.sql.DataSource;

import org.locationtech.jts.geom.Geometry;

import com.google.common.base.Joiner;
import com.google.common.base.Strings;
import com.google.common.io.CharStreams;

import nz.org.riskscape.engine.RiskscapeException;
import nz.org.riskscape.engine.Tuple;
import nz.org.riskscape.engine.output.AxisSwapper;
import nz.org.riskscape.engine.output.RiskscapeWriter;
import nz.org.riskscape.engine.output.StructFlattener;
import nz.org.riskscape.engine.output.StructFlattener.StructMapping;
import nz.org.riskscape.engine.problem.ProblemFactory;
import nz.org.riskscape.engine.problem.SeverityLevel;
import nz.org.riskscape.engine.types.Geom;
import nz.org.riskscape.engine.types.Struct;
import nz.org.riskscape.problem.Problem;
import nz.org.riskscape.problem.ProblemException;
import nz.org.riskscape.problem.Problems;
import nz.org.riskscape.problem.ResultOrProblems;

import lombok.Data;
import lombok.Getter;
import lombok.RequiredArgsConstructor;

/**
 * Base class for output formats that write to a JDBC database.
 */
@RequiredArgsConstructor
public abstract class BaseJdbcOutputStore  {

  /**
   * The default number of insert statements that will be committed as a batch by the {@link RiskscapeWriter}.
   */
  public static final int DEFAULT_BATCH_INSERT_SIZE = 1000;

  public interface LocalProblems extends ProblemFactory {

    /**
     * Problem to give some context to other problems that can occur.
     * @param when    when the problem occurred. Expected to be either 'create' or 'write'
     * @param output  URI representing the output store
     * @return
     */
    Problem failedWhen(String when, URI output);

    /**
     * Use when appending to an existing table is not possible because the existing table has a different
     * schema than required.
     */
    Problem cannotAppendTableStructureMismatch(String tableName, String requiredTable, String existingTable);

    /**
     * A hint to add as a child to {@link #cannotAppendTableStructureMismatch(String, String, String) }
     */
    @SeverityLevel(Problem.Severity.INFO)
    Problem cannotAppendTableHint();
  }

  public static final LocalProblems PROBLEMS = Problems.get(LocalProblems.class);

  /**
   * Data to allow the mapping of a {@link StructMapping} to a DB column.
   */
  @Data
  @RequiredArgsConstructor
  public static class StructMappingToColumnMapping {

    public StructMappingToColumnMapping(StructMapping structMapping, String dbColumnType,
        Function<Object, Object> mapper) {
      this.structMapping = structMapping;
      this.dbColumnType = dbColumnType;
      this.mapper = mapper;
      this.sqlType = Optional.empty();
    }

    /**
     * The {@link StructMapping} that is the source of the data.
     */
    private final StructMapping structMapping;

    /**
     * The SQL type that will be required for storing the mapped values.
     */
    private final String dbColumnType;

    /**
     * A {@link Function} to translate the RiskScape value into one that is a type compatible with
     * {@link #getDbColumnType()}.
     */
    private final Function<Object, Object> mapper;

    /**
     * An optional integer that when supplied will be passed to
     * {@link PreparedStatement#setObject(int, java.lang.Object, int) } in argument three.
     *
     * If not given the {@link PreparedStatement#setObject(int, java.lang.Object) } is used instead.
     */
    private final Optional<Integer> sqlType;
  }

  public class JdbcRiskscapeWriter extends RiskscapeWriter {

    @Getter
    private final String tableName;
    private final List<StructMappingToColumnMapping> mappings;
    private final Connection conn;
    private final boolean closeConnection;
    private final PreparedStatement insertStmt;
    @Getter
    private final URI storedAt;
    private final AtomicInteger tupleCounter = new AtomicInteger();

    /**
     * Create a {@link RiskscapeWriter} that will write {@link Tuple}s to the given table.
     *
     * @param tableName the name of the table to be written to
     * @param mappings  struct to column mappings to the given table. There must be an entry for every column
     *                  in tableName and these must be in order.
     * @param conn      connection to use when inserting records to the DB. This connection will be closed when
     *                  {@link RiskscapeWriter#close() } is called.
     * @param closeConnection when true {@link #conn} will be closed when this writer is closed. Allows for
     *                        cases where the connection is shared between writers, but in this case something
     *                        else will need to close the connection at the appropriate time
     * @param storedAt  a URI representing where the table being written to. This will be returned when
     *                  {@link RiskscapeWriter#getStoredAt() } is called.
     * @throws SQLException if the insert statement cannot be built
     */
    JdbcRiskscapeWriter(String tableName, List<StructMappingToColumnMapping> mappings, Connection conn,
          boolean closeConnection, URI storedAt)
        throws SQLException {
      this.tableName = tableName;
      this.mappings = mappings;
      this.conn = conn;
      this.closeConnection = closeConnection;
      this.storedAt = storedAt;

      this.insertStmt = conn.prepareStatement(String.format("insert into %s values(%s)",
          quoteIdentifier(tableName, conn.getMetaData()),
          Joiner.on(",").join(Collections.nCopies(mappings.size(), "?"))
      ));
    }

    @Override
    public void write(Tuple tuple) {
      try {
        for (int i = 0; i < mappings.size(); i++) {
          StructMappingToColumnMapping mapping = mappings.get(i);
          Object value = mapping.structMapping.getAccessExpression().evaluate(tuple);
          if (value != null) {
            value = mapping.mapper.apply(value);
          }
          // PreparedStatement parameter indexes start at 1, hence the + 1
          if (mapping.sqlType.isPresent()) {
            insertStmt.setObject(i + 1, value, mapping.sqlType.get());
          } else {
            insertStmt.setObject(i + 1, value);
          }
        }
        insertStmt.executeUpdate();

        if (tupleCounter.incrementAndGet() % DEFAULT_BATCH_INSERT_SIZE == 0) {
          // we have reached the batch commit size, time to commit
          synchronized (conn) {
            conn.commit();
          }
        }
      } catch (SQLException e) {
        throw new RiskscapeException(PROBLEMS.failedWhen("write", storedAt).withChildren(Problems.caught(e)));
      }
    }

    @Override
    public void close() throws IOException {
      try {
        synchronized (conn) {
          insertStmt.close();
          conn.commit();
          if (closeConnection) {
            conn.close();
          }
        }
      } catch (SQLException e) {
        throw new IOException(e);
      }
    }

  }

  protected final DataSource ds;
  protected final boolean replaceExistingTables;

  /**
   * Get a writer that can write the given {@link Struct}.
   *
   * @param type the type that the writer will write
   * @param name the desired table name to write tuples to
   * @return
   */
  public ResultOrProblems<JdbcRiskscapeWriter> writerFor(Struct type, String name) {
    return doWriterFor(type, name, false);
  }

  /**
   * Return a writer for the given type and name that will append to existing tables should they
   * exist and have the same shape.
   *
   * If the table already exists but with a different shape (schema) then
   * {@link LocalProblems#cannotAppendTableStructureMismatch(String, String, String) } will result.
   *
   * Note, that {@link #replaceExistingTables} has no effect when appending.
   */
  public ResultOrProblems<JdbcRiskscapeWriter> appendingWriterFor(Struct type, String name) {
    return doWriterFor(type, name, true);
  }

  protected ResultOrProblems<JdbcRiskscapeWriter> doWriterFor(Struct type, String name, boolean appending) {
    // We use the struct flattener to flatten out the struct
    StructFlattener flattener = new StructFlattener();
    List<StructMapping> flattened = flattener.flatten(type, new StructFlattener.LastMemberNamer());

    // now we look for a geometry. if one is found that member will be a spatial entry.
    StructMapping firstGeomMapping = flattened.stream()
        .filter(m -> m.getType().findAllowNull(Geom.class).isPresent())
        .findFirst()
        .orElse(null);

    List<Problem> warnings = new ArrayList<>();
    Connection conn = null;
    try {
      // we synchronize to prevent any race conditions whilst creating the layer. this in particularly
      // important for implementations that use a shared connection (GeoPackage) but also generally as
      // multiple outputs may otherwise attempt to make the same modifications to common tables, such as
      // attempting duplicate inserts to the spatial ref table.
      synchronized (this) {
        // Get a connection for the writer. Note that this connection is not closed here except when
        // exceptions are caught. On normal use the writer is responsible for closing the connection.
        conn = getConnection();

        // set auto commit to false. This allow us to rollback all the table creation DDL if there are
        // any problems so should prevent any corrupt DB's in that case.
        conn.setAutoCommit(false);

        DatabaseMetaData metaData = conn.getMetaData();
        if (name.contains(metaData.getIdentifierQuoteString())) {
          // if the desired table name contains the quote character we replace it early with a harmless
          // underscore. We do this now for the table name so that the code that is checking for existing
          // tables will be checking for the correct table. Otherwise the table name could change slightly
          // when the create table calls quoteIdentifier() on the tableName.
          name = name.replaceAll(metaData.getIdentifierQuoteString(), "_");
        }

        // The implementation may require other tables to exist, such as a GeoPackage which needs
        // a contents and spatial reference tables
        initDbIfNecessary(conn);

        List<StructMappingToColumnMapping> toColumnMappings = new ArrayList<>();
        for (StructMapping sm: flattened) {
          toColumnMappings.add(
              toColumnMapping(sm, firstGeomMapping, conn)
                  .drainWarnings(warning -> warnings.add(warning))
                  .getOrThrow()
          );
        }

        List<String> existingTables = getExistingTableNames(metaData);
        String tableName = name;
        if (existingTables.contains(tableName)) {
          if (replaceExistingTables && ! appending) {
            // the desired table name already exists, but user says replace it
            // note that replace is mutally exclusive with appending
            deleteTable(name, conn);
          } else {
            if (appending) {
              try {
                checkExistingTablesMatches(tableName, metaData, toColumnMappings);
                // the existing table has the same schema we can return a writer to append to it.
                return ResultOrProblems.of(createJdbcWriter(
                    tableName, toColumnMappings, conn, getTableURI(urlEncode(tableName))
                ), warnings);
              } catch (ProblemException pe) {
                // no dice, the existing table does not have the required structure.
                // we could just create a new table but we'd get a new table on every run which isn't great.
                // we could look through the tables with trailing indexes and use the first match but that
                // is getting hard to reason which table will be selected. So until a use-case exists for
                // something else we just bail. (Note appending is only used for manifiest writing).
                return pe.toResult();
              }
            }
            // we add an index to the end of the table name and increment it until we find
            // one that does not exist.
            int counter = 1;
            while (existingTables.contains(tableName)) {
              tableName = String.format("%s_%d", name, counter++);
            }
          }
        }

        StringBuilder sb = new StringBuilder(String.format("create table %s(", quoteIdentifier(tableName, metaData)));
        for (int i = 0; i < toColumnMappings.size(); i++) {
          if (i > 0) {
            sb.append(", ");
          }
          StructMappingToColumnMapping mapping = toColumnMappings.get(i);
          sb.append(String.format("%s %s",
              quoteIdentifier(mapping.structMapping.getKey(), metaData),
              mapping.dbColumnType)
          );
        }
        sb.append(")");
        String createTableSql = sb.toString();

        try (Statement s = conn.createStatement()) {
          s.execute(createTableSql);
          postTableCreate(tableName, firstGeomMapping, conn);
          // Let's commit the table creation DDL before returning the writer
          conn.commit();
          return ResultOrProblems.of(createJdbcWriter(
              tableName, toColumnMappings, conn, getTableURI(urlEncode(tableName))
          ), warnings);
        }
      }
    } catch (ProblemException | SQLException e) {
      // we catch the ProblemException (rather than handle with ProblemException.catching so we can
      // rollback and close the DB connection as well.
      try {
        if (conn != null) {
          conn.rollback();
          conn.close();
        }
      } catch (SQLException ee) {}
      ResultOrProblems<JdbcRiskscapeWriter> result;
      if (e instanceof ProblemException) {
        result = ((ProblemException) e).toResult();
      } else {
        result = ResultOrProblems.failed(Problems.caught(e));
      }
      // we compose the problem to give some context on when the problem occurred and with what. We use
      // getTableURI has that should give the best context to the user.
      return result.composeProblems(PROBLEMS.failedWhen("create", getTableURI(urlEncode(name))));
    }
  }

  /**
   * Check if the existing table has the same structure as the requiredMappings.
   */
  private void checkExistingTablesMatches(String tableName, DatabaseMetaData metadata,
      List<StructMappingToColumnMapping> requiredMappings) throws ProblemException, SQLException {
    ResultSet rs = metadata.getColumns(null, null, tableName, null);
    // we are going to collect the existing columns as we go, just in case they are needed for error reporting.
    List<String> existingColumns = new ArrayList<>();
    boolean matches = true;
    int i = 0;
    while(rs.next()) {
      String existingName = rs.getString("COLUMN_NAME");
      String existingType = rs.getString("TYPE_NAME");
      if ("_text".equals(existingType) && rs.getInt("DATA_TYPE") == java.sql.Types.ARRAY) {
        // PostGIS hack for supporting TEXT[]. The array types don't get returned in the metadata the same
        // as the simple types, so we make it work for this exception.
        existingType = "TEXT[]";
      }
      existingColumns.add(String.format("%s %s", existingName, existingType));
      if (i >= requiredMappings.size()) {
        // there are extra columns in the table so we don't match. but we need to keep running the loop
        // to collect all the existing columns.
        matches = false;
        continue;
      }
      StructMappingToColumnMapping requiredColumn = requiredMappings.get(i++);
      if (!requiredColumn.structMapping.getKey().equals(existingName)) {
        // column name mismatch
        matches = false;
      }
      if (!requiredColumn.dbColumnType.equalsIgnoreCase(existingType)) {
        // column type mismatch
        matches = false;
      }
    }
    if (requiredMappings.size() != existingColumns.size()) {
      // there is a different number of columns
      matches = false;
    }
    if (! matches) {
      String requiredTable = requiredMappings.stream()
          .map(mapping -> String.format("%s %s", mapping.structMapping.getKey(), mapping.dbColumnType))
          .collect(Collectors.joining(", ", tableName + "(", ")"));
      String existingTable = existingColumns.stream()
          .collect(Collectors.joining(", ", tableName + "(", ")"));
      throw new ProblemException(PROBLEMS.cannotAppendTableStructureMismatch(tableName, requiredTable, existingTable)
          .withChildren(PROBLEMS.cannotAppendTableHint()));
    }
  }

  /**
   * Get a {@link Connection} that will be used to both create and then populate the table.
   *
   * Note that if this method is overridden to return a shared connection then
   * {@link #createJdbcWriter(java.lang.String, java.util.List, java.sql.Connection, java.net.URI) } should
   * also be overridden to return a {@link JdbcRiskscapeWriter} that will not close the shared connection.
   * In this case the implementation will need to take other steps to ensure that the connection is closed.
   */
  protected Connection getConnection() throws SQLException {
    return ds.getConnection();
  }

  protected JdbcRiskscapeWriter createJdbcWriter(String tableName, List<StructMappingToColumnMapping> mappings,
      Connection conn, URI storedAt) throws SQLException {
    return new JdbcRiskscapeWriter(tableName, mappings, conn, replaceExistingTables, storedAt);
  }

  protected String urlEncode(String value) {
    try {
      return URLEncoder.encode(value, "UTF-8");
    } catch (UnsupportedEncodingException e) {
      throw new RuntimeException(e); // utf-8 is pretty fundamental
    }
  }

  /**
   * Quotes a database identifier (table or column name) with the database's quote character, replacing any quote
   * characters that might appear within the identifier itself with an underscore (`_`)
   */
  protected String quoteIdentifier(String identifier, DatabaseMetaData metaData) throws SQLException {
    // NB Java 9 has a method for this, d'oh
    String quoteChar = metaData.getIdentifierQuoteString();
    if (identifier.contains(quoteChar)) {
      identifier = identifier.replaceAll(quoteChar, "_");
    }
    return String.format("%s%s%s", quoteChar, identifier, quoteChar);
  }

  /**
   * Get a URI that represents the given table at the JDBC data source. Will be returned from
   * {@link RiskscapeWriter#getStoredAt() } when writing is complete.
   *
   * @param urlEncodedtableName the table that is being written with URL encoding
   * @return URI representing the table
   */
  public abstract URI getTableURI(String urlEncodedtableName);

  /**
   * Allows implementations to initialize the DB with any format required tables should this be
   * necessary.
   */
  protected void initDbIfNecessary(Connection conn) throws SQLException {}

  /**
   * Provide a {@link StructMappingToColumnMapping} for the given structMapping.
   *
   * firstGeomMapping is also provided to allow implementations that only allow one spatial entry to
   * map subsequent geometries, maybe to text WKT.
   *
   * @param structMapping the struct mapping to map
   * @param firstGeomMapping the first geometry mapping, if one exists
   * @return a mapper
   */
  protected abstract ResultOrProblems<StructMappingToColumnMapping> toColumnMapping(StructMapping structMapping,
      StructMapping firstGeomMapping, Connection conn) throws ProblemException;

  /**
   * Allow implementations to perform any other actions that may be required now that a feature
   * table has been created. This could be to register the table such as a GeoPackage will register
   * a table in 'gpkg_contents' and 'gpkg_geometry_columns' for a geometry column.
   *
   * @param tableName the name of the table that has been created
   * @param firstGeomMapping the first geometry mapping from that table, or null if none exist
   * @param conn a connection that can be used for further sql statements.
   */
  protected void postTableCreate(String tableName, StructMapping firstGeomMapping, Connection conn)
      throws SQLException, ProblemException {
  }

  /**
   * Get a list of names of the existing tables and views that exist in the DB.
   * @param metaData DB metadata to find existing tables in
   * @return list of existing table and view names
   * @throws SQLException if database metadata cannot be obtained
   */
  protected List<String> getExistingTableNames(DatabaseMetaData metaData) throws SQLException {
    List<String> tableNames = new ArrayList<>();
    try (ResultSet found = metaData.getTables(null, null, "", new String[]{"TABLE", "VIEW"})) {
      while (found.next()) {
        tableNames.add(found.getString("TABLE_NAME"));
      }
    }
    return tableNames;
  }

  protected void deleteTable(String tableName, Connection conn) throws SQLException {
    try (Statement stmt = conn.createStatement()) {
      stmt.executeUpdate("drop table " + quoteIdentifier(tableName, conn.getMetaData()));
    }
  }

  /**
   * Run the sql script content from the given {@link InputStream}.
   *
   * SQL statements may span multiple lines but must be terminated by a ';'. Comment lines starting
   * with '--' are allowed.
   */
  protected void runScript(InputStream is, Connection conn) throws IOException, SQLException {
    try (BufferedReader reader = new BufferedReader(new InputStreamReader(is));
        Statement s = conn.createStatement()) {
      StringBuilder sqlBuilder = new StringBuilder();
      for (String line: CharStreams.readLines(reader)) {
        line = line.trim();
        if (!(Strings.isNullOrEmpty(line) || line.startsWith("--"))) {
          sqlBuilder.append(line).append(" ");
          if (line.endsWith(";")) {
            s.addBatch(sqlBuilder.toString());
            // reset the builder ready for the next statement
            sqlBuilder.setLength(0);
          }
        }
      }
      s.executeBatch();
    }
  }

  /**
   * Prepare a {@link Geometry} for writing, which means to swap the axis if necessary and set the writeSRID.
   *
   * These steps are performed on a copy of the input geometry to ensure that is it not altered as that
   * could have bad affects should it be used in later RiskScape processing.
   *
   * @param geom        geometry to prepare
   * @param writeSRID   The SRID that the prepared geometry should have
   * @param axisSwapper Axis swapper if required
   * @return  the prepared geometry
   */
  protected Geometry prepareGeometry(Geometry geom, int writeSRID, Optional<AxisSwapper> axisSwapper) {
    // The axis swapper, if present, will make a copy of the geom, but if the srid needs to be set then we must clone
    // the geometry - we can't assume we're the last save step in the pipeline that'll see the given geometry
    if (axisSwapper.isPresent()) {
      geom = axisSwapper.get().swapAxis(geom);
    } else if (writeSRID != geom.getSRID()) {
      geom = geom.copy();
    }

    // this is a noop if the srid hasn't actually changed and saves an extra test
    geom.setSRID(writeSRID);

    return geom;
  }

}
