package au.gov.amsa.gt;
   
   import java.io.File;
   import java.io.IOException;
   import java.io.InputStream;
   import java.io.Serializable;
   import java.io.Writer;
   import java.nio.file.Files;
   import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.concurrent.atomic.AtomicInteger;
  
  import org.geotools.data.DataStore;
  import org.geotools.data.DataStoreFinder;
  import org.geotools.data.shapefile.ShapefileDataStore;
  import org.geotools.data.simple.SimpleFeatureCollection;
  import org.geotools.data.simple.SimpleFeatureIterator;
  import org.geotools.data.simple.SimpleFeatureSource;
  import org.geotools.feature.simple.SimpleFeatureBuilder;
  import org.geotools.feature.simple.SimpleFeatureTypeBuilder;
  import org.geotools.geojson.feature.FeatureJSON;
  import org.geotools.geometry.jts.JTS;
  import org.geotools.referencing.CRS;
  import org.opengis.feature.simple.SimpleFeature;
  import org.opengis.referencing.FactoryException;
  import org.opengis.referencing.crs.CoordinateReferenceSystem;
  import org.opengis.referencing.operation.MathTransform;
  
  import com.vividsolutions.jts.awt.PointShapeFactory.X;
  import com.vividsolutions.jts.geom.Coordinate;
  import com.vividsolutions.jts.geom.Geometry;
  import com.vividsolutions.jts.geom.GeometryFactory;
  import com.vividsolutions.jts.geom.prep.PreparedGeometry;
  import com.vividsolutions.jts.geom.prep.PreparedGeometryFactory;
  
  public final class Shapefile {
  
      private final DataStore datastore;
      private final static GeometryFactory GEOMETRY_FACTORY = new GeometryFactory();
      // 0 = not loaded, not closed, 1 = loaded, 2 = closed
      private final AtomicInteger state = new AtomicInteger(0);
      private static final int NOT_LOADED = 0;
      private static final int LOADED = 1;
      private static final int CLOSED = 2;
  
      private volatile List<PreparedGeometry> geometries;
      private final double bufferDistance;
      private volatile List<SimpleFeature> features = new ArrayList<SimpleFeature>();
      private volatile CoordinateReferenceSystem crs;
  
      private Shapefile(File file, double bufferDistance) {
          this.bufferDistance = bufferDistance;
          try {
              datastore = new ShapefileDataStore(file.toURI().toURL());
          } catch (IOException e) {
              throw new RuntimeException(e);
          }
      }
  
      public static Shapefile from(File file) {
          return from(file, 0);
      }
  
      public static Shapefile from(File file, double bufferDistance) {
          return new Shapefile(file, bufferDistance);
      }
  
      public static void createPolygon(List<Coordinate> coords, File output) {
          ShapefileCreator.createPolygon(coords, output);
      }
  
      public static Shapefile fromZip(InputStream is) {
          return fromZip(is, 0);
      }
  
      public static DataStore fromZipAsDataStore(InputStream is) {
          try {
              File directory = Files.createTempDirectory("shape-").toFile();
              ZipUtil.unzip(is, directory);
              Map<String, Serializable> map = new HashMap<>();
              map.put("url", directory.toURI().toURL());
              return DataStoreFinder.getDataStore(map);
          } catch (IOException e) {
              throw new RuntimeException(e);
          }
      }
  
      public static Shapefile fromZip(InputStream is, double bufferDistance) {
          try {
              File directory = Files.createTempDirectory("shape-").toFile();
              ZipUtil.unzip(is, directory);
              File f = directory.listFiles(x -> x.getName().endsWith(".shp"))[0];
              return new Shapefile(f, bufferDistance);
          } catch (IOException e) {
              throw new RuntimeException(e);
          }
     }
 
     private Shapefile load() {
         if (state.compareAndSet(NOT_LOADED, LOADED)) {
             try {
                 final List<PreparedGeometry> geometries = new ArrayList<>();
                 for (String typeName : datastore.getTypeNames()) {
                     SimpleFeatureSource source = datastore.getFeatureSource(typeName);
                     crs = source.getBounds().getCoordinateReferenceSystem();
                     final SimpleFeatureCollection features = source.getFeatures();
                     SimpleFeatureIterator it = features.features();
                     while (it.hasNext()) {
                         SimpleFeature feature = it.next();
                         Geometry g = (Geometry) feature.getDefaultGeometry();
                         if (bufferDistance > 0)
                             g = g.buffer(bufferDistance);
                         geometries.add(PreparedGeometryFactory.prepare(g));
                         this.features.add(feature);
                     }
                     it.close();
                 }
                 this.geometries = geometries;
             } catch (IOException e) {
                 throw new RuntimeException(e);
             }
         } else if (state.get() == CLOSED)
             throw new RuntimeException("Shapefile is closed and can't be accessed");
         return this;
     }
 
     public List<PreparedGeometry> geometries() {
         load();
         return geometries;
     }
 
     public boolean contains(double lat, double lon) {
         load();
         return GeometryUtil.contains(GEOMETRY_FACTORY, geometries, lat, lon);
     }
 
     public Rect mbr() {
         // TODO assumes that shapefile is using WGS84?
         load();
         Rect r = null;
         for (PreparedGeometry g : geometries) {
             Coordinate[] v = g.getGeometry().getEnvelope().getCoordinates();
             System.out.println(Arrays.toString(v));
             Rect rect = new Rect(v[0].y, v[0].x, v[2].y, v[2].x);
             if (r == null)
                 r = rect;
             else
                 r = r.add(rect);
         }
         return r;
     }
 
     public void close() {
         if (state.compareAndSet(NOT_LOADED, CLOSED) || state.compareAndSet(LOADED, CLOSED))
             datastore.dispose();
     }
 
     public void writeGeoJson(Writer writer, String targetCrsName) {
         load();
 
         try {
             CoordinateReferenceSystem targetCrs = CRS.decode(targetCrsName);
             MathTransform transform = CRS.findMathTransform(crs, targetCrs);
             FeatureJSON f = new FeatureJSON();
             f.writeCRS(targetCrs, writer);
 
             features.stream().forEach(feature -> {
                 try {
                     SimpleFeatureTypeBuilder tb = new SimpleFeatureTypeBuilder();
                     tb.setCRS(targetCrs);
                     SimpleFeatureBuilder b = new SimpleFeatureBuilder(feature.getFeatureType());
                     b.add(JTS.transform((Geometry) feature.getDefaultGeometry(), transform));
                     f.writeFeature(b.buildFeature(feature.getID()), writer);
                 } catch (Exception e) {
                     throw new RuntimeException(e);
                 }
             });
         } catch (IOException | FactoryException e) {
             throw new RuntimeException(e);
         }
     }
 }