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

import java.io.File;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.URI;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.UUID;
import java.util.function.Consumer;
import java.util.function.Function;

import javax.xml.stream.XMLEventFactory;
import javax.xml.stream.XMLEventReader;
import javax.xml.stream.XMLEventWriter;
import javax.xml.stream.XMLInputFactory;
import javax.xml.stream.XMLOutputFactory;
import javax.xml.stream.XMLStreamException;
import javax.xml.stream.events.XMLEvent;

import org.geotools.geometry.jts.ReferencedEnvelope;
import org.geotools.referencing.wkt.Formattable;
import org.geotools.api.referencing.FactoryException;
import org.geotools.api.referencing.crs.CoordinateReferenceSystem;
import org.geotools.api.referencing.operation.TransformException;

import com.google.common.collect.ImmutableMap;

import lombok.Getter;
import lombok.RequiredArgsConstructor;
import nz.org.riskscape.engine.SRIDSet;
import nz.org.riskscape.engine.data.Bookmark;
import nz.org.riskscape.engine.data.BookmarkResolvers;
import nz.org.riskscape.engine.output.FileSystemPipelineOutputStore.FileSystemPipelineContainer;
import nz.org.riskscape.engine.pipeline.ExecutionContext;
import nz.org.riskscape.engine.pipeline.sink.SaveSink;
import nz.org.riskscape.engine.problem.ProblemFactory;
import nz.org.riskscape.engine.relation.Relation;
import nz.org.riskscape.engine.resource.CreateException;
import nz.org.riskscape.problem.Problem;
import nz.org.riskscape.problem.Problem.Severity;
import nz.org.riskscape.problem.Problems;
import nz.org.riskscape.problem.ResultOrProblems;

/**
 * Writes a QGIS project file (*.qgs) from as a collection of the output layers
 */
public class QGSWriter {

  /**
   * The default output location to write the project file to.
   */
  private static final String TEMPLATE_FILE = "/nz/org/riskscape/engine/qgis/template.xml";

  /**
   * QGIS project file generated by the writer
   */
  private static final String FILE = "project.qgs";

  /**
   * Bookmark options we should use when trying to parse the output file
   */
  private static final Map<String, List<?>> BOOKMARK_OPTIONS = ImmutableMap.of(
      // skip any geometry warnings (as well as the expensive work of verifying geometry)
      "validate-geometry", Arrays.asList("off")
  );

  public static final Consumer<FileSystemPipelineContainer> WRITE_PROJECT_FILE = c -> write(c);

  /**
   * Uses the given {@link FileSystemPipelineContainer} to save the QGIS project file
   */
  public static void write(FileSystemPipelineContainer container) {
    try {
      QGSWriter writer = new QGSWriter(container.getExecutionContext(), container.getSinks());
      container.writeFile(FILE, "application/x-qgis-project", os -> writer.write(os));
    } catch (Exception e) {
      container.getExecutionContext().getProblemSink().log(
          PROBLEMS.errorWritingFile().withSeverity(Severity.WARNING)
              .withChildren(Problems.caught(e)));
    }
  }

  public interface LocalProblems extends ProblemFactory {
    Problem errorWritingFile();
  }

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

  private XMLEventWriter writer;
  private final Map<String, SaveSink> sinks;
  private final XMLEventFactory eventFactory = XMLEventFactory.newInstance();
  @Getter
  private final List<Layer> layers = new ArrayList<Layer>();

  @RequiredArgsConstructor
  static class Layer {
    @Getter private final String name;
    @Getter private final ReferencedEnvelope bounds;
    @Getter private final String source;
    private final UUID uuid = UUID.randomUUID();

    public String getId() {
      return (name + "_" + uuid.toString()).replace("-", "_");
    }

    public CoordinateReferenceSystem getCrs() {
      return bounds.getCoordinateReferenceSystem();
    }
  }

  public QGSWriter(ExecutionContext context, Map<String, SaveSink> sinks) {
    this.sinks = sinks;
    generateLayers(context);
  }

  // constructor for testing so we can feed in Relations directly
  QGSWriter(Map<String, SaveSink> sinks, Function<URI, ResultOrProblems<Relation>> getRelation) {
    this.sinks = sinks;
    generateLayers(getRelation);
  }

  /**
   * Populate the layers attribute with Layer objects containing the
   * relevant subset of data from sinks
   */
  private void generateLayers(ExecutionContext executionContext) {
    BookmarkResolvers resolvers = executionContext.getEngine().getBookmarkResolvers();
    generateLayers(uri -> resolvers.getData(Bookmark.fromURI(uri)
          .addUnparsed(BOOKMARK_OPTIONS), executionContext, Relation.class));
  }

  private void generateLayers(Function<URI, ResultOrProblems<Relation>> getRelation) {
    for (Entry<String, SaveSink> entry : sinks.entrySet()) {
      String name = entry.getKey();
      URI uri = entry.getValue().getStoredAt();
      String source = new File(uri).toString();

      // read the relation data that's been written to the URI
      getRelation.apply(uri).ifPresent(relation -> {
          ReferencedEnvelope bounds = relation.calculateBounds()
              .orElse(new ReferencedEnvelope(SRIDSet.EPSG4326_LONLAT));

          layers.add(new Layer(name, bounds, source));
      });
    }
    // present the layers in alphabetical order in QGIS
    layers.sort((a, b) -> a.getName().compareTo(b.getName()));
  }

  ReferencedEnvelope getBounds() {
    return getLayers().stream()
        .map(l -> l.getBounds())
        .map(b -> {
          if (b.getCoordinateReferenceSystem().equals(SRIDSet.EPSG4326_LONLAT)) {
            return b;
          }
          try {
            return b.transform(SRIDSet.EPSG4326_LONLAT, false);
          } catch (TransformException | FactoryException e) {}
          return new ReferencedEnvelope();
        })
        .filter(b -> !b.isNull())
        .reduce(
            new ReferencedEnvelope(SRIDSet.EPSG4326_LONLAT),
            (bounds, layerBounds) -> {
              bounds.expandToInclude(layerBounds);
              return bounds;
            }
        );
  }

  /**
   * Initialise a XMLEventReader instance that reads from the QGS template resource.
   */
  private XMLEventReader createReader(InputStream is) throws XMLStreamException {
    return XMLInputFactory.newInstance().createXMLEventReader(is);
  }

  /**
   * Initialise a XMLEventWriter instance that writes to the given output stream.
   */
  private XMLEventWriter createWriter(OutputStream os) throws XMLStreamException {
    return XMLOutputFactory.newInstance().createXMLEventWriter(os);
  }

  private void writeString(String s) throws XMLStreamException {
    writer.add(eventFactory.createCharacters(s));
  }

  /**
   * Write a XML start element with value and optional attributes.
   */
  private void writeStartElement(String name, String value, String... attributes) throws XMLStreamException {

    writer.add(eventFactory.createStartElement("", "", name));

    for (int i = 0; i + 1 < attributes.length; i += 2) {
      writer.add(eventFactory.createAttribute(attributes[i], attributes[i+1]));
    }

    if (!value.isEmpty()) {
      writeString(value);
    }
  }

  /**
   * Write a XML end element
   */
  private void writeEndElement(String name) throws XMLStreamException {
    writer.add(eventFactory.createEndElement("", "", name));
  }

  /** Write a XML element */
  private void writeElement(String name, String value, String... attributes) throws XMLStreamException {
    writeStartElement(name, value, attributes);
    writer.add(eventFactory.createEndElement("", "", name));
  }

  private void writeLayerOrderLayers() throws XMLStreamException {
    for (Layer layer : getLayers()) {
      writeElement("layer", "", "id", layer.getId());
    }
  }

  private void writeMapLayers() throws XMLStreamException {

    for (Layer layer : getLayers()) {
      writeStartElement("maplayer", "", "type", "vector");
      writeElement("id", layer.getId());
      writeElement("datasource", layer.getSource());
      writeElement("layername", layer.getName());
      writeStartElement("srs", "");
      writeStartElement("spatialrefsys", "");

      CoordinateReferenceSystem crs = layer.getCrs();
      if (crs != null) {
        Formattable formattableCRS = (Formattable) crs;
        String crsWKT = formattableCRS.toWKT(0).replace("\n", "");

        writeElement("wkt", crsWKT);
      } else {
        writeElement("wkt", "");
      }

      writeEndElement("spatialrefsys");
      writeEndElement("srs");
      writeEndElement("maplayer");
    }
  }

  private void writeLayerTreeLayers() throws XMLStreamException {
    for (Layer layer : getLayers()) {
      writeElement("layer-tree-layer", "",
      "name", layer.getName(), "source", layer.getSource(), "id", layer.getId());
    }
  }

  private void writeExtent() throws XMLStreamException {
    ReferencedEnvelope bounds = getBounds();
    String xmin = "";
    String ymin = "";
    String xmax = "";
    String ymax = "";

    if (!bounds.isNull()) {
      xmin = String.valueOf(bounds.getMinX());
      ymin = String.valueOf(bounds.getMinY());
      xmax = String.valueOf(bounds.getMaxX());
      ymax = String.valueOf(bounds.getMaxY());
    }

    writeStartElement("extent", "");
    writeElement("xmin", xmin);
    writeElement("ymin", ymin);
    writeElement("xmax", xmax);
    writeElement("ymax", ymax);
    writeEndElement("extent");
  }

  /**
   * Writes the project file to the given output stream.
   */
  private void writeAll(OutputStream os) throws XMLStreamException {
    InputStream is = getClass().getResourceAsStream(TEMPLATE_FILE);
    XMLEventReader reader = createReader(is);
    writer = createWriter(os);

    // parse template and write a modified copy
    while (reader.hasNext()) {
      XMLEvent event = reader.nextEvent();
      writer.add(event);

      if (event.isStartElement()) {
        // write new elements inside this element
        switch (event.asStartElement().getName().getLocalPart()) {
          case "layerorder":
            writeLayerOrderLayers();
            break;
          case "projectlayers":
            writeMapLayers();
            break;
          case "mapcanvas":
            writeExtent();
            break;
          default:
            break;
        }
      } else if (event.isEndElement()) {
        // write new elements after this element
        switch (event.asEndElement().getName().getLocalPart()) {
          case "customproperties":
            writeLayerTreeLayers();
            break;
          default:
            break;
        }
      }
    }

    writer.flush();
    writer.close();
    reader.close();
  }

  /**
   * Writes the project file to the given output stream.
   */
  public void write(OutputStream os) throws CreateException {
    try {
      writeAll(os);
    } catch (XMLStreamException e) {
      throw new CreateException(Problems.caught(e));
    }
  }
}
