package org.simantics.maps.elevation.server;

import java.awt.geom.Rectangle2D;
import java.io.Closeable;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import org.geotools.geometry.DirectPosition2D;
import org.geotools.geometry.Envelope2D;
import org.geotools.geometry.jts.ReferencedEnvelope;
import org.geotools.referencing.CRS;
import org.opengis.geometry.DirectPosition;
import org.opengis.referencing.crs.CoordinateReferenceSystem;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.github.benmanes.caffeine.cache.Caffeine;
import com.github.benmanes.caffeine.cache.LoadingCache;
import com.vividsolutions.jts.geom.Coordinate;
import com.vividsolutions.jts.geom.Envelope;
import com.vividsolutions.jts.index.strtree.STRtree;

public class TiffTileInterface implements Closeable {

    private static final Logger LOGGER = LoggerFactory.getLogger(TiffTileInterface.class);

    private Path tilesFolder;
    private Map<Path, Envelope> envelopes = new ConcurrentHashMap<>();
    private LoadingCache<Path, TiffInterface> interfaceCache;
    private int openInterfacesSize;
    private STRtree index;

    public TiffTileInterface(Path tilesFolder) {
        this(tilesFolder, 5);
    }

    public TiffTileInterface(Path tilesFolder, int openInterfacesSize) {
        if (!Files.isDirectory(tilesFolder)) {
            throw new IllegalArgumentException("tilesFolder has to be a folder: " + tilesFolder.toAbsolutePath());
        }
        this.tilesFolder = tilesFolder;
        this.index = new STRtree();
        this.openInterfacesSize = openInterfacesSize;

        this.interfaceCache = Caffeine.newBuilder()
            .maximumSize(this.openInterfacesSize)
            .removalListener((key, gdalInterface, cause) -> ((TiffInterface) gdalInterface).close())
            .build(key -> new TiffInterface(key));
        
        try {
            initializeIndex();
        } catch (IOException e) {
            LOGGER.error("Could not initialize index for folder {}", tilesFolder, e);
        }
    }

    private TiffInterface openTifInterface(Path tifFile) {
        return interfaceCache.get(tifFile);
    }

    private Stream<Path> allTiffFiles() throws IOException {
        return Files.walk(tilesFolder).filter(Files::isRegularFile).filter(tif -> tif.getFileName().toString().endsWith(".tif"));
    }

    public void initializeIndex() throws IOException {
        LOGGER.info("Initializing index..");
        AtomicInteger counter = new AtomicInteger();
        allTiffFiles().parallel().forEach(tifFile -> {
            TiffInterface tifInterface = openTifInterface(tifFile);
            Envelope2D coords = tifInterface.getCornerCoords();
            try {
                ReferencedEnvelope refEnv = new ReferencedEnvelope(coords);
                ReferencedEnvelope targetEnv = refEnv.transform(c4326, false, 30);

                synchronized(index) {
                    index.insert(targetEnv, tifFile);
                }
                envelopes.put(tifFile, targetEnv);
            } catch (Exception e) {
                LOGGER.error("Could not initialize index for file {}", tifFile, e);
            } finally {
                tifInterface.close();
                int current = counter.getAndIncrement();
                if (current % 100 == 0) {
                    LOGGER.info("    {}", current);
                }
            }
        });
    }

    public Collection<Rectangle2D> getBoundingBoxes() {
        Collection<Rectangle2D> rects = envelopes.values().stream().map(env -> {
            double x = env.getMinX();
            double y = env.getMinY();
            return new Rectangle2D.Double(x, y, env.getMaxX() - x, env.getMaxY() - y);
        }).collect(Collectors.toList());
        return rects;
    }

    private static CoordinateReferenceSystem c4326;

    static {
        try {
            c4326 = CRS.decode("EPSG:4326");
        } catch (Exception e) {
            LOGGER.error("Could not initialize EPSG:4326", e);
        }
    }

    public Number lookup(double x, double y) {
        LOGGER.trace("Looking up x={} y={}", x, y);
        DirectPosition p = new DirectPosition2D(c4326, x, y);
        List<Path> tifFiles = (List<Path>) index.query(new Envelope(new Coordinate(x, y)));
        for (Path tifFile : tifFiles) {
            TiffInterface tifInterface = openTifInterface(tifFile);
            if (tifInterface.contains(p)) {
                try {
                    return tifInterface.lookup(p);
                } finally {
                    tifInterface.close();
                }
            } else {
                //System.out.println("not found");
            }
        }
        return new Double(0); // use 0 by default for now
    }

    @Override
    public void close() throws IOException {
        interfaceCache.invalidateAll();
        interfaceCache.cleanUp();
        
        envelopes.clear();
        envelopes = null;
        index = null;
        interfaceCache = null;
    }
}