/*
 * 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.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.net.URI;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.time.Instant;
import java.time.format.DateTimeFormatter;
import java.util.Date;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.function.Function;

import javax.sql.DataSource;

import org.apache.commons.dbcp.BasicDataSource;
import org.geotools.geopkg.geom.GeoPkgGeomWriter;
import org.geotools.referencing.CRS;
import org.geotools.referencing.crs.DefaultEngineeringCRS;
import org.locationtech.jts.geom.Geometry;
import org.geotools.api.referencing.FactoryException;
import org.geotools.api.referencing.crs.CoordinateReferenceSystem;

import lombok.Getter;
import nz.org.riskscape.engine.GeometryProblems;
import nz.org.riskscape.engine.OutputProblems;
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.GeoPackagePipelineOutputStore;
import nz.org.riskscape.engine.output.RiskscapeWriter;
import nz.org.riskscape.engine.output.StructFlattener;
import nz.org.riskscape.engine.types.Referenced;
import nz.org.riskscape.engine.types.Struct;
import nz.org.riskscape.problem.ProblemException;
import nz.org.riskscape.problem.ProblemSink;
import nz.org.riskscape.problem.Problems;
import nz.org.riskscape.problem.ResultOrProblems;

/**
 * For writing to GeoPackage files.
 *
 * Note that writing to GeoPackage files uses a shared connection. It is important that
 * {@link #close() is called to close this shared connection.
 */
public class GeoPackageOutputStore extends BaseJdbcOutputStore implements AutoCloseable {

  /**
   * Get a {@link RiskscapeWriter} that can write the given type/name to geoPackageFile.
   *
   * Note that the GeoPackageOutputStore will be closed when the returned RiskscapeWriter is
   * closed.
   *
   * This method is only appropriate when a single layer is to be written to the GeoPackage. If multiple
   * layers are to be written then {@link #writerFor(Struct, String) } should be used.
   *
   * @param geoPackageFile  the GeoPackage file to be written to
   * @param type  the type to write
   * @param name  the desired name to write to
   * @param replaceExistingTables true if existing tables should be replaced
   * @return riskscape writer or problems preventing one from being created
   */
  public static ResultOrProblems<RiskscapeWriter> writerFor(File geoPackageFile,
      Struct type, String name, boolean replaceExistingTables) {
    GeoPackageOutputStore store = new GeoPackageOutputStore(geoPackageFile, false);
    return store.writerFor(type, name)
        // we need to map the returned writer to one that will close the store. This is because
        // the store has the connection to share between all writers that use that store. But from
        // here, there will only every be one writer.
        .map(writer -> new RiskscapeWriter() {

      @Override
      public void write(Tuple value) {
        writer.write(value);
      }

      @Override
      public URI getStoredAt() {
        return geoPackageFile.toURI();
      }

      @Override
      public void close() throws IOException {
        // first close the writer, it will likely have transactions that need committing.
        writer.close();
        try {
          store.close();
        } catch (Exception e) {
          throw new IOException(e);
        }
      }

    });

  }

  private static DataSource createDataSource(File geoPackageFile) {
    BasicDataSource ds = new BasicDataSource();
    ds.setDriverClassName("org.sqlite.JDBC");
    ds.setUrl(String.format("jdbc:sqlite:%s", geoPackageFile.getAbsolutePath()));
    ds.setDefaultReadOnly(false);

    return ds;
  }

  @Getter
  private final File geoPackageFile;
  // the connection that will be used by this output store. the connection is shared by all of the writers
  // that use this store.
  private Connection connection = null;

  public GeoPackageOutputStore(File geoPackageFile, boolean replaceExistingTables) {
    super(createDataSource(geoPackageFile), replaceExistingTables);
    this.geoPackageFile = geoPackageFile;
  }

  @Override
  protected Connection getConnection() throws SQLException {
    // the connection needs to be shared by all the outputs that use it
    if (this.connection == null) {
      this.connection = super.getConnection();
    }
    return this.connection;
  }

  @Override
  protected JdbcRiskscapeWriter createJdbcWriter(String tableName, List<StructMappingToColumnMapping> mappings,
      Connection conn, URI storedAt) throws SQLException {
    // we need to create a JdbcRiskscapeWriter that will not close the shared connection
    return new JdbcRiskscapeWriter(tableName, mappings, conn, false, storedAt);
  }

  @Override
  public void close() throws Exception {
    if (connection != null) {
      connection.close();
    }
  }

  @Override
  protected ResultOrProblems<StructMappingToColumnMapping> toColumnMapping(StructFlattener.StructMapping structMapping,
      StructFlattener.StructMapping firstGeomMapping, Connection conn) throws ProblemException {
    Class type = structMapping.getType().internalType();
    if (Geometry.class.isAssignableFrom(type)) {
      if (structMapping != firstGeomMapping) {
        // we have found a subsequent geometry. we'll map this to text
        return ResultOrProblems.of(new StructMappingToColumnMapping(structMapping, "text", g -> g.toString()));
      }
      Referenced referenced = structMapping.getType().findAllowNull(Referenced.class)
          .orElseThrow(() -> new ProblemException(GeometryProblems.get().notReferenced(structMapping.getType())));
      Optional<AxisSwapper> axisSwapper = AxisSwapper.getForceXY(referenced.getCrs(), null, ProblemSink.DEVNULL);

      // GeoPackage mandates a XY axis order. So we'll use the AxisSwapper if needed to do the switch axis thing
      CoordinateReferenceSystem crs = axisSwapper.map(AxisSwapper::getNewCrs)
          .orElse(referenced.getCrs());

      // The SRID is also included in the GeoPackage geometry encoding. But this should be the srs_id value
      // that is used in the GeoPackages gpkg_spatial_ref table.
      int geoPkgSRID = toSridCode(crs, conn);

      GeoPkgGeomWriter geomWriter = new GeoPkgGeomWriter();
      ByteArrayOutputStream bos = new ByteArrayOutputStream();

      return ResultOrProblems.of(new StructMappingToColumnMapping(structMapping, type.getSimpleName(), g -> {
        Geometry geom = (Geometry) g;
        try {
          // we call reset to clear out any previous geometries that have been written to bos,
          bos.reset();
          geomWriter.write(prepareGeometry(geom, geoPkgSRID, axisSwapper), bos);
          return bos.toByteArray();
        } catch (IOException e) {
          throw new RiskscapeException(Problems.caught(e));
        }
      }));
    }
    if (type == Boolean.class) {
      return ResultOrProblems.of(new StructMappingToColumnMapping(structMapping, "BOOLEAN",
          b -> ((Boolean) b) ? 1 : 0));
    }
    if (type == Long.class) {
      return ResultOrProblems.of(new StructMappingToColumnMapping(structMapping, "INTEGER", Function.identity()));
    }
    if (type == Double.class) {
      return ResultOrProblems.of(new StructMappingToColumnMapping(structMapping, "REAL", Function.identity()));
    }
    if (type == Date.class) {
      return ResultOrProblems.of(new StructMappingToColumnMapping(structMapping, "TEXT",
          d -> DateTimeFormatter.ISO_INSTANT.format(Instant.ofEpochMilli(((Date) d).getTime()))));
    }
    if (type == String.class) {
      return ResultOrProblems.of(new StructMappingToColumnMapping(structMapping, "TEXT", Function.identity()));
    }
    // we map any unhandled types to string. this is in keeping with what the Shapefile writer does and
    // is required to allow types that aren't normally expected to be saved (eg coverage type) to be
    // saved in raw-results without causing an error.
    return ResultOrProblems.of(new StructMappingToColumnMapping(structMapping, "TEXT", x -> Objects.toString(x)),
        OutputProblems.get().outputTypeAsText(
            structMapping.getKey(),
            structMapping.getType(),
            new GeoPackagePipelineOutputStore())
      );
  }

  @Override
  protected void postTableCreate(String tableName, StructFlattener.StructMapping geomMapping, Connection conn)
      throws SQLException, ProblemException {
    if (geomMapping == null) {
      try (PreparedStatement stmt = conn.prepareStatement(
          "insert into gpkg_contents(table_name, data_type, identifier) values(?, 'attributes', ?)")) {
        stmt.setString(1, tableName);
        stmt.setString(2, tableName);
        stmt.executeUpdate();
      }
    } else {
      CoordinateReferenceSystem crs = geomMapping.getType().findAllowNull(Referenced.class)
          .map(Referenced::getCrs)
          .orElse(null);
      int geoPackageSrid = toSridCode(crs, conn);
      try (PreparedStatement stmt = conn.prepareStatement(
          "insert into gpkg_contents(table_name, data_type, identifier, srs_id) values(?, 'features', ?, ?)")) {
        stmt.setString(1, tableName);
        stmt.setString(2, tableName);
        stmt.setInt(3, geoPackageSrid);
        stmt.executeUpdate();
      }
      try (PreparedStatement stmt = conn.prepareStatement(
          "insert into gpkg_geometry_columns values(?, ?, ?, ?, 0, 0)")) {
        stmt.setString(1, tableName);
        stmt.setString(2, geomMapping.getKey());
        stmt.setString(3, geomMapping.getType().internalType().getSimpleName());
        stmt.setInt(4, geoPackageSrid);
        stmt.executeUpdate();
      }
    }
  }

  @Override
  public URI getTableURI(String tableName) {
    // we create a URI that is compatible with the GeoPackage resolver to read the layer
    return URI.create(String.format("%s?layer=%s", geoPackageFile.toURI(), tableName, "UTF-8"));
  }

  /**
   * Initializes the GeoPackage with required tables should they not exist already
   * @param conn
   */
  @Override
  protected void initDbIfNecessary(Connection conn) {
    try {
      if (! getExistingTableNames(conn.getMetaData()).contains("gpkg_contents")) {
        // if gpkg_contents does not exist as a table or view then we assume this ds has not yet been initialized
        // with the required geopackage tables. Lets add them now.
        InputStream initStream = getClass().getResourceAsStream(
            String.format("/%s/geopackage-init.sql", getClass().getPackage().getName().replaceAll("\\.", "/"))
        );
        runScript(initStream, conn);
      }
    } catch (IOException | SQLException e) {
      throw new RiskscapeException(Problems.caught(e));
    }
  }

  @Override
  protected void deleteTable(String tableName, Connection conn) throws SQLException {
    // The GeoPacakge format has a bunch of tables that are used to register layers in.
    // These typically have unique constraints etc. But either way the old ones should be purged before
    // we add new entries.
    List<String> existingTables = getExistingTableNames(conn.getMetaData());
    try (Statement stmt = conn.createStatement()) {
      stmt.execute("delete from gpkg_contents where table_name = '" + tableName + "'");
      stmt.execute("delete from gpkg_geometry_columns where table_name = '" + tableName + "'");
      // gpkg_extensions is an optional table
      if (existingTables.contains("gpkg_extensions")) {
        stmt.execute("delete from gpkg_extensions where table_name = '" + tableName + "'");
      }

    }
    // now we call super so it can delete the target table
    super.deleteTable(tableName, conn);
  }

  /**
   * Convert the CRS to an SRID code. Registering it into the 'gpkg_spatial_ref_sys' table if
   * necessary.
   *
   * @param crs the CRS to get the SRID code of
   * @return the SRID code for given CRS
   */
  private int toSridCode(CoordinateReferenceSystem crs, Connection conn) throws ProblemException {
    try (Statement stmt = conn.createStatement()) {
      Integer epsgCode = null;

      try {
        epsgCode = CRS.lookupEpsgCode(crs, false);
      } catch (FactoryException ex) {
        // ths CRS doesn't look like an EPSG CRS. We'll need to handle it some other way but lets
        // defer that until we know there is a use case for non EPSG CRSs.
        throw new ProblemException(OutputProblems.get().unsupportedCrs(crs));
      }
      if (epsgCode == null) {
        // ths CRS doesn't look like an EPSG CRS. We'll need to handle it some other way but lets
        // defer that until we know there is a use case for non EPSG CRSs. GL729 asks is supporting
        // non EPSG CRS's is required
        if (DefaultEngineeringCRS.class.isAssignableFrom(crs.getClass())) {
          // GeoPackage defined SRID for Undefined cartesian SRS
          return -1;
        }
        throw new ProblemException(OutputProblems.get().unsupportedCrs(crs));
      }
      ResultSet rs = stmt.executeQuery(String.format(
          "select srs_id from gpkg_spatial_ref_sys where organization_coordsys_id = %d and organization like 'EPSG'",
          epsgCode
      ));
      if (rs.next()) {
        // the CRS is already registered
        return rs.getInt("srs_id");
      }
      // The CRS needs to be registered
      try (PreparedStatement insertStmt =
          conn.prepareStatement("insert into gpkg_spatial_ref_sys values(?, ?, 'EPSG', ?, ?, ?)")) {
        insertStmt.setString(1, "EPSG:" + epsgCode);
        insertStmt.setInt(2, epsgCode);
        insertStmt.setInt(3, epsgCode);
        insertStmt.setString(4, crs.toWKT());
        insertStmt.setString(4, "EPSG:" + epsgCode);
        insertStmt.execute();
      }
      return epsgCode;
    } catch (SQLException e) {
      throw new RiskscapeException(Problems.caught(e));
    }
  }

}
