/*
 * 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.IOException;
import java.io.OutputStreamWriter;
import java.io.Writer;
import java.net.URI;
import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;
import java.util.stream.Collectors;

import javax.xml.stream.XMLOutputFactory;
import javax.xml.stream.XMLStreamException;
import javax.xml.stream.XMLStreamWriter;

import org.locationtech.jts.geom.Coordinate;
import org.locationtech.jts.geom.Geometry;
import org.locationtech.jts.geom.GeometryCollection;
import org.locationtech.jts.geom.LineString;
import org.locationtech.jts.geom.LinearRing;
import org.locationtech.jts.geom.Point;
import org.locationtech.jts.geom.Polygon;

import com.google.common.collect.Lists;

import nz.org.riskscape.engine.Project;
import nz.org.riskscape.engine.RiskscapeException;
import nz.org.riskscape.engine.SRIDSet;
import nz.org.riskscape.engine.Tuple;
import nz.org.riskscape.engine.projection.FlattenProjection;
import nz.org.riskscape.engine.projection.Projector;
import nz.org.riskscape.engine.query.TupleUtils;
import nz.org.riskscape.engine.resource.CreateHandle;
import nz.org.riskscape.engine.types.Nullable;
import nz.org.riskscape.engine.types.Struct;
import nz.org.riskscape.engine.types.Struct.StructMember;
import nz.org.riskscape.engine.types.Type;
import nz.org.riskscape.engine.types.Types;
import nz.org.riskscape.engine.util.Pair;
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 nz.org.riskscape.problem.StandardCodes;
import lombok.Getter;

/**
 * A {@link RiskscapeWriter} that writes KML (Keyhole Markup Language)
 */
public class KmlWriter extends RiskscapeWriter {

  public static final String KML_NAMESPACE = "http://www.opengis.net/kml/2.2";
  /**
   * These are the standard KML data fields that could potentially be set from RiskScape data. They
   * are string fields except for 'visibility' and 'open' which are boolean.
   */
  public static final List<String> KML_DATA_MEMBERS = Lists.newArrayList("name", "visibility", "open", "address",
      "phoneNumber", "description");

  public static ResultOrProblems<KmlWriter> of(Struct type, CreateHandle handle, Project project) {
    return ProblemException.catching(() -> {
      FlattenProjection flattener = new FlattenProjection(".");
      Projector flattenProjector = flattener.getProjectionFunction(type).getOrThrow();

      StructMember geomMember = TupleUtils.findGeometryMember(
          flattenProjector.getProjectedType(), TupleUtils.FindOption.OPTIONAL);
      if (geomMember == null) {
        throw new ProblemException(Problem.error(StandardCodes.GEOMETRY_REQUIRED, type));
      }
      List<StructMember> dataMembers = new ArrayList<>();
      List<Pair<String, StructMember>> extendedDataMembers = new ArrayList<>();

      List<String> usedMemberNames = new ArrayList<>(KML_DATA_MEMBERS);
      usedMemberNames.add("Geometry");  // what the geometry member is called when read.
      List<StructMember> members = new ArrayList<>(flattenProjector.getProjectedType().getMembers());
      members.remove(geomMember);
      for (String dataMember : KML_DATA_MEMBERS) {
        StructMember member = findMember(dataMember, members);
        if (member != null) {
          if ((dataMember.equals("visibility") || dataMember.equals("open"))
              && member.getType().getUnwrappedType() != Types.BOOLEAN) {
            // we can only use visibility or open for boolean types. otherwise they go down to the extended
            // types section
            continue;
          }
          dataMembers.add(member);
          members.remove(member);
        }
      }
      for (StructMember member : members) {
        int i = 1;
        String name = member.getKey();
        while (usedMemberNames.contains(name)) {
          name = String.format("%s-%d", member.getKey(), i++);
        }
        extendedDataMembers.add(Pair.of(name, member));
      }

      return new KmlWriter(flattenProjector, geomMember, dataMembers, extendedDataMembers,
          handle, project.getSridSet());
    });
  }

  private static StructMember findMember(String named, List<StructMember> members) {
    return members.stream()
        .filter(member -> member.getKey().equals(named))
        .findFirst()
        .orElse(null);
  }

  private final Projector flattenProjector;
  private final StructMember geometryMember;
  private final List<StructMember> dataMembers;
  private final List<Pair<String, StructMember>> extendedDataMembers;

  private final CreateHandle handle;
  private final SRIDSet sridSet;
  @Getter
  private URI storedAt = null;
  private Writer baseWriter;
  private XMLStreamWriter xmlWriter = null;

  private KmlWriter(Projector flattenProjector, StructMember geometryMember, List<StructMember> dataMembers,
      List<Pair<String, StructMember>> extendedDataMembers, CreateHandle handle, SRIDSet sridSet) {
    this.flattenProjector = flattenProjector;
    this.geometryMember = geometryMember;
    this.dataMembers = dataMembers;
    this.extendedDataMembers = extendedDataMembers;
    this.handle = handle;
    this.sridSet = sridSet;

    initWriter();
  }

  @Override
  public void write(Tuple input) {
    Tuple flattened = flattenProjector.apply(input);

    try {
      xmlWriter.writeStartElement("Placemark");
      for (StructMember dataMember : dataMembers) {
        Object value = flattened.fetch(dataMember);
        if (value == null) {
          // don't try to write a null value at all.
          continue;
        }
        xmlWriter.writeStartElement(dataMember.getKey());
        xmlWriter.writeCharacters(Objects.toString(value));
        xmlWriter.writeEndElement();
      }
      if (!extendedDataMembers.isEmpty()) {
        xmlWriter.writeStartElement("ExtendedData");
        xmlWriter.writeStartElement("SchemaData");
        xmlWriter.writeAttribute("schemaUrl", "#riskscape");

        for (Pair<String, StructMember> extendedMember : extendedDataMembers) {
          StructMember member = extendedMember.getRight();
          Object value = flattened.fetch(member);
          if (value == null) {
            // don't try to write a null value at all.
            continue;
          }
          if (member == geometryMember) {
            // kml geometries are always in lon/lat (CRS84)
            value = sridSet.reproject((Geometry) value, sridSet.get(SRIDSet.EPSG4326_LONLAT));
          }
          xmlWriter.writeStartElement("SimpleData");
          xmlWriter.writeAttribute("name", extendedMember.getLeft());
          xmlWriter.writeCharacters(Objects.toString(value));
          xmlWriter.writeEndElement();
        }
        xmlWriter.writeEndElement(); // SchemaData
        xmlWriter.writeEndElement(); // ExtendedData
      }

      // Now we write the geometry entry which has to be in the KML mandated CRS:84.
      Geometry geom = flattened.fetch(geometryMember);
      if (geom != null) {
        // only non null geometries are written
        writeGeometry(sridSet.reproject(geom, sridSet.get(SRIDSet.EPSG4326_LONLAT)));
      }

      xmlWriter.writeEndElement();
    } catch (XMLStreamException e) {
      e.printStackTrace();
      throw new RiskscapeException(Problems.caught(e));
    }
  }

  /**
   * Initializes the xml writer. We don't do this until the first tuple is written so there is no
   * file created for no tuples.
   */
  private void initWriter() {
    try {
      baseWriter = new OutputStreamWriter(handle.getOutputStream(), Charset.forName("UTF-8"));
      XMLOutputFactory xmlOutputFactory = XMLOutputFactory.newInstance();
      xmlOutputFactory.setProperty(XMLOutputFactory.IS_REPAIRING_NAMESPACES, true);
      xmlWriter = xmlOutputFactory.createXMLStreamWriter(baseWriter);
      xmlWriter.writeStartDocument("UTF-8", "1.0");
      xmlWriter.setDefaultNamespace(KML_NAMESPACE);
      xmlWriter.writeStartElement("kml");
      xmlWriter.writeStartElement("Document");
      xmlWriter.writeAttribute("id", "root_doc");
      if (!extendedDataMembers.isEmpty()) {
        xmlWriter.writeStartElement("Schema");
        xmlWriter.writeAttribute("name", "riskscape");
        xmlWriter.writeAttribute("id", "riskscape");
        for (Pair<String, StructMember> extendedMember : extendedDataMembers) {
          StructMember member = extendedMember.getRight();
          if (member == geometryMember) {
            continue;
          }
          // KML has a limited set of allowed types.
          // See https://developers.google.com/kml/documentation/kmlreference#schema
          // This block maps our types to the KML types and falls back to string if there is no better match.
          Type memberType = Nullable.strip(member.getType().getUnwrappedType());
          String type = "string";
          if (memberType == Types.INTEGER) {
            type = "int";
          } else if (memberType == Types.FLOATING) {
            type = "double";
          } else if (memberType == Types.BOOLEAN) {
            type = "bool";
          }
          xmlWriter.writeStartElement("SimpleField");
          xmlWriter.writeAttribute("name", extendedMember.getLeft());
          xmlWriter.writeAttribute("type", type);
          xmlWriter.writeEndElement();
        }
        xmlWriter.writeEndElement();
      }
      xmlWriter.writeStartElement("Folder");
    } catch (Exception e) {
      throw new RiskscapeException(Problems.caught(e));
    }
  }

  @Override
  public void close() throws IOException {
    if (xmlWriter != null) {
      try {
        xmlWriter.writeEndElement();  // </kml:Folder>
        xmlWriter.writeEndElement();  // </kml:Document>
        xmlWriter.writeEndElement();  // </kml:kml>
        xmlWriter.writeEndDocument();
        xmlWriter.close();
        baseWriter.close();
      } catch (XMLStreamException e) {
        throw new RiskscapeException(Problems.caught(e));
      }
      storedAt = handle.store();
    }
  }

  private void writeGeometry(Geometry g) throws XMLStreamException {
    if (g instanceof Point) {
      writePoint((Point) g);
    } else if (g instanceof LinearRing) {
      writeLinearRing((LinearRing) g);
    } else if (g instanceof LineString) {
      writeLineString((LineString) g);
    } else if (g instanceof Polygon) {
      writePolygon((Polygon) g);
    } else if (g instanceof GeometryCollection) {
      writeGeometryCollection((GeometryCollection) g);
    } else {
      throw new IllegalArgumentException("Geometry type not supported: " + g.getGeometryType());
    }
  }

  private void writePoint(Point point) throws XMLStreamException {
    xmlWriter.writeStartElement("Point");
    writeCoordinates(point.getCoordinates());
    xmlWriter.writeEndElement();
  }

  private void writeLineString(LineString line) throws XMLStreamException {
    xmlWriter.writeStartElement("LineString");
    writeCoordinates(line.getCoordinates());
    xmlWriter.writeEndElement();
  }

  private void writeLinearRing(LinearRing ring) throws XMLStreamException {
    xmlWriter.writeStartElement("LinearRing");
    writeCoordinates(ring.getCoordinates());
    xmlWriter.writeEndElement();
  }

  private void writePolygon(Polygon poly) throws XMLStreamException {
    xmlWriter.writeStartElement("Polygon");
    xmlWriter.writeStartElement("outerBoundaryIs");
    writeLinearRing(poly.getExteriorRing());
    xmlWriter.writeEndElement();
    for (int i = 0; i < poly.getNumInteriorRing(); i++) {
      xmlWriter.writeStartElement("innerBoundaryIs");
      writeLinearRing(poly.getInteriorRingN(i));
      xmlWriter.writeEndElement();
    }
    xmlWriter.writeEndElement();
  }

  private void writeGeometryCollection(GeometryCollection multi) throws XMLStreamException {
    xmlWriter.writeStartElement("MultiGeometry");
    for (int i = 0; i < multi.getNumGeometries(); i++) {
      writeGeometry(multi.getGeometryN(i));
    }
    xmlWriter.writeEndElement();
  }

  private void writeCoordinates(Coordinate[] coordinates) throws XMLStreamException {
    xmlWriter.writeStartElement("coordinates");
    // KML coordinates come in in space separated tuples of lon,lat[,alt]
    // TODO consider if we should include alt (if present) and maybe we can reduce precision
    String coordinateString = Arrays.stream(coordinates)
        .map(c -> String.format("%f,%f", c.getX(), c.getY()))
        .collect(Collectors.joining(" "));
    xmlWriter.writeCharacters(coordinateString);
    xmlWriter.writeEndElement();
  }

}
