/*
 * 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.postgis;

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.util.ArrayList;
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.geotools.referencing.CRS;
import org.locationtech.jts.geom.Geometry;
import org.locationtech.jts.io.WKBWriter;
import org.geotools.api.referencing.FactoryException;
import org.geotools.api.referencing.crs.CoordinateReferenceSystem;
import org.postgresql.ds.PGSimpleDataSource;

import lombok.Getter;
import lombok.NonNull;
import lombok.extern.slf4j.Slf4j;

import nz.org.riskscape.engine.GeometryProblems;
import nz.org.riskscape.engine.OutputProblems;
import nz.org.riskscape.engine.RiskscapeException;
import nz.org.riskscape.engine.defaults.data.jdbc.BaseJdbcOutputStore;
import nz.org.riskscape.engine.output.AxisSwapper;
import nz.org.riskscape.engine.output.BaseJdbcPipelineOutputContainer;
import nz.org.riskscape.engine.output.Format;
import nz.org.riskscape.engine.output.PipelineOutputContainer;
import nz.org.riskscape.engine.output.PipelineOutputOptions;
import nz.org.riskscape.engine.output.PipelineOutputStore;
import nz.org.riskscape.engine.output.SinkParameters;
import nz.org.riskscape.engine.output.StructFlattener;
import nz.org.riskscape.engine.pipeline.RealizedPipeline;
import nz.org.riskscape.engine.pipeline.sink.SaveSink;
import nz.org.riskscape.engine.types.RSList;
import nz.org.riskscape.engine.types.Referenced;
import nz.org.riskscape.engine.types.Text;
import nz.org.riskscape.problem.Problem;
import nz.org.riskscape.problem.ProblemException;
import nz.org.riskscape.problem.ProblemSink;
import nz.org.riskscape.problem.Problems;
import nz.org.riskscape.problem.ResultOrProblems;

/**
 * A {@link PipelineOutputStore} that will write results to a PostGIS database.
 */
@Slf4j
public class PostGISPipelineOutputStore implements PipelineOutputStore {

  private static final PostGISProblems POSTGIS_PROBLEMS = PostGISProblems.get();

  /**
   * Convenience to get {@link PostGISProblems#connectionFailure(java.net.URI, java.lang.String) }
   * along with the child problem tips to help the user resolve the connection failure.
   */
  public static Problem toConnectionFailure(URI location, Exception cause) {
    return POSTGIS_PROBLEMS.connectionFailure(location, cause.getMessage())
        .withChildren(
            POSTGIS_PROBLEMS.connectionTipSpelling(),
            POSTGIS_PROBLEMS.connectionTipPassword(),
            POSTGIS_PROBLEMS.connectionTipConnectivity()
        );
  }

  @Getter
  private final String id = "postgis";


  class PostGISOutputContainer extends BaseJdbcPipelineOutputContainer<PostGISJdbcAdaptor> {

    private URI postGisUri;

    PostGISOutputContainer(RealizedPipeline pipeline, URI postGisUri, PipelineOutputOptions options, DataSource ds) {
      super(
        PostGISPipelineOutputStore.this,
        new PostGISJdbcAdaptor(postGisUri, ds, options.isReplace()),
        pipeline,
        options
      );
      this.postGisUri = postGisUri;
    }

    @Override
    public URI getStoredAt() {
      // rebuild a 'sanitized' non-jdbc version of the url
      StringBuilder builder = new StringBuilder("postgis://");

      builder.append(postGisUri.getHost());

      if (postGisUri.getPort() > 0) {
        builder.append(":").append(postGisUri.getPort());
      }

      if (postGisUri.getRawPath() != null) {
        builder.append(postGisUri.getRawSchemeSpecificPart());
      }

      return URI.create(builder.toString());
    }

    @Override
    protected ResultOrProblems<SaveSink> createSink(SinkParameters sinkParameters) {
      List<Problem> warnings = new ArrayList<>();
      if (sinkParameters.getFormat().isPresent()) {
        warnings.add(OutputProblems.get().userSpecifiedFormatIgnored(sinkParameters.getFormat().get(), getStore()));
      }

      return super.createSink(sinkParameters).withMoreProblems(warnings);
    }
  }

  class PostGISJdbcAdaptor extends BaseJdbcOutputStore {

    private final URI location;

    PostGISJdbcAdaptor(URI location, javax.sql.DataSource ds, boolean replaceExistingTables) {
      super(ds, replaceExistingTables);
      this.location = location;
    }

    @Override
    public URI getTableURI(String tableName) {
      return URI.create(String.format("%s?layer=%s", location, tableName));
    }

    @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)) {
        Referenced referenced = structMapping.getType().findAllowNull(Referenced.class)
            .orElseThrow(() -> new ProblemException(GeometryProblems.get().notReferenced(structMapping.getType())));

        // PostGIS stores geometry in the CRS defined order.
        Optional<AxisSwapper> axisSwapper = AxisSwapper.getForceXY(referenced.getCrs(), null, ProblemSink.DEVNULL);
        CoordinateReferenceSystem crs = axisSwapper.map(AxisSwapper::getNewCrs)
            .orElse(referenced.getCrs());

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

        WKBWriter geomWriter = new WKBWriter(2, true);

        return ResultOrProblems.of(new StructMappingToColumnMapping(structMapping,
            String.format("geometry(%s, %d)", type.getSimpleName(), postgisSRID),
            g -> {
              Geometry geom = (Geometry) g;
              return geomWriter.write(prepareGeometry(geom, postgisSRID, axisSwapper));
            }));

      }
      if (type == Boolean.class) {
        return ResultOrProblems.of(new StructMappingToColumnMapping(structMapping, "BOOLEAN", Function.identity()));
      }
      if (type == Long.class) {
        return ResultOrProblems.of(new StructMappingToColumnMapping(structMapping, "BIGINT", Function.identity()));
      }
      if (type == Double.class) {
        return ResultOrProblems.of(new StructMappingToColumnMapping(structMapping, "DOUBLE PRECISION",
            Function.identity()));
      }
      if (type == Date.class) {
        return ResultOrProblems.of(new StructMappingToColumnMapping(structMapping, "TIMESTAMP", Function.identity(),
            Optional.of(java.sql.Types.TIMESTAMP)));
      }
      if (type == String.class) {
        return ResultOrProblems.of(new StructMappingToColumnMapping(structMapping, "TEXT", Function.identity()));
      }
      if (structMapping.getType().findAllowNull(RSList.class)
          .map(listType -> listType.find(Text.class)).isPresent()) {
        // mapping text lists is added to support writing the manifest output tables
        return ResultOrProblems.of(new StructMappingToColumnMapping(structMapping, "TEXT[]", list -> {
          List values = (List)list;
          return values.toArray(new String[values.size()]);
        }));
      }
      // 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(),
            PostGISPipelineOutputStore.this));
    }

    @Override
    protected void postTableCreate(String tableName, StructFlattener.StructMapping firstGeomMapping, Connection conn)
        throws SQLException, ProblemException {
      // TODO consider creating a spatial index
    }

    /**
     * Convert the CRS to an SRID code. Registering it into the '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()) {
        if (!getExistingTableNames(conn.getMetaData()).contains("spatial_ref_sys")) {
          // the PostGIS extension adds the spatial_ref_sys table to a database for storing the CRS's.
          // if that table is missing then the postgis extension has not been added to the database
          // and geometry cannot be stored.
          throw new ProblemException(POSTGIS_PROBLEMS.databaseNotSpatial());
        }
        Integer epsgCode = null;
        try {
          epsgCode = CRS.lookupEpsgCode(crs, false);

        } catch (FactoryException ex) {
          log.error("EPSG code lookup failed - crs will not be supported for postgis output", ex);
          // something is very wrong with the CRS EPSG lookup. We'll let the if null handling below
          // return a problem rather than copy pasting that here.
        }
        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.
          throw new ProblemException(OutputProblems.get().unsupportedCrs(crs));
        }
        ResultSet rs = stmt.executeQuery(String.format(
            "select srid from spatial_ref_sys where auth_srid = %d and auth_name like 'EPSG'",
            epsgCode
        ));
        if (rs.next()) {
          // the CRS is already registered
          return rs.getInt("srid");
        }
        // The CRS needs to be registered
        try (PreparedStatement insertStmt
            = conn.prepareStatement("insert into spatial_ref_sys values(?, ?, ?, ?)")) {
          // FIXME there's no reason why the srid (primary key) has as to match the auth_id abd we are naively assuming
          // that the primary key will be 'unoccupied' - we probably need to do a max
          // lookup or something here, or maybe have a start-from id that we use to avoid collisions.  hmmm...
          // See https://postgis.net/docs/using_postgis_dbmanagement.html#spatial_ref_sys for more details
          insertStmt.setInt(1, epsgCode);
          insertStmt.setString(2, "EPSG");
          insertStmt.setInt(3, epsgCode);
          insertStmt.setString(4, crs.toWKT());
          insertStmt.execute();
        }
        return epsgCode;
      } catch (SQLException e) {
        throw new RiskscapeException(Problems.caught(e));
      }
    }
  }

  @Override
  public int isApplicable(@NonNull URI outputLocation) {
    return PostGISConnectionOptions.URI_SCHEME.equals(outputLocation.getScheme()) ? PRIORITY_HIGH : PRIORITY_NA;
  }

  @Override
  public ResultOrProblems<PipelineOutputContainer> create(URI outputLocation, RealizedPipeline pipeline,
      PipelineOutputOptions options) {
    List<Problem> warnings = new ArrayList<>();

    Format userSpecifiedDefault = options.getFormat().orElse(null);
    if (userSpecifiedDefault != null) {
      if (! "postgis".equals(userSpecifiedDefault.getId())) {
        warnings.add(OutputProblems.get().userSpecifiedFormatIgnored(userSpecifiedDefault, this));
      }
      // set the default to null, we're ignoring it so don't want warnings every time
      // it is ignored
      options.setFormat(null);
    }
    if (options.isChecksum()) {
      warnings.add(OutputProblems.get().checksumNotSupported(this));
    }

    PostGISConnectionOptions cnxOptions = new PostGISConnectionOptions();
    cnxOptions.setLocation(outputLocation);
    cnxOptions.setEnvironment(System.getenv());

    return cnxOptions.validate()
        .drainWarnings(warning -> warnings.add(warning))
        .composeProblems((severity, problems) ->
            // we compose the problems to give the user some handy context to help them figure out
            // what is wrong with the location
            POSTGIS_PROBLEMS.invalidLocation(outputLocation)
                .withSeverity(severity)
                .withChildren(problems)
        )
        .flatMap(connectionOptions -> {
          PGSimpleDataSource ds = new PGSimpleDataSource();
          ds.setUrl(connectionOptions.toJdbcUri());

          try (Connection conn = ds.getConnection()) {
            conn.getMetaData();
          } catch (SQLException e) {
            return ResultOrProblems.failed(toConnectionFailure(outputLocation, e));
          }

          return ResultOrProblems.of(new PostGISOutputContainer(pipeline, outputLocation, options, ds), warnings);
        });
  }
}
