package com.lingniu.ingest.sink.archive;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.URI;
import java.nio.file.Files;
import java.nio.file.LinkOption;
import java.nio.file.Path;
import java.nio.file.StandardOpenOption;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

public final class LocalArchiveStore implements ArchiveStore {

    private final Path root;
    private final ConcurrentMap<Path, Object> appendLocks = new ConcurrentHashMap<>();

    public LocalArchiveStore(String root) {
        this(rootPath(root));
    }

    public LocalArchiveStore(Path root) {
        this.root = root.toAbsolutePath().normalize();
    }

    @Override
    public String put(String key, InputStream data, long length) throws IOException {
        if (data == null) {
            throw new IllegalArgumentException("data must not be null");
        }
        Path target = resolve(key);
        createDirectoriesInsideRoot(target.getParent());
        rejectSymlinkPath(target);
        try (OutputStream out = Files.newOutputStream(
                target,
                StandardOpenOption.CREATE,
                StandardOpenOption.TRUNCATE_EXISTING,
                StandardOpenOption.WRITE,
                LinkOption.NOFOLLOW_LINKS)) {
            data.transferTo(out);
        }
        return target.toUri().toString();
    }

    @Override
    public String append(String key, byte[] chunk) throws IOException {
        Path target = resolve(key);
        synchronized (appendLocks.computeIfAbsent(target, ignored -> new Object())) {
            createDirectoriesInsideRoot(target.getParent());
            rejectSymlinkPath(target);
            Files.write(
                    target,
                    chunk == null ? new byte[0] : chunk,
                    StandardOpenOption.CREATE,
                    StandardOpenOption.APPEND,
                    StandardOpenOption.WRITE,
                    LinkOption.NOFOLLOW_LINKS);
        }
        return target.toUri().toString();
    }

    @Override
    public InputStream get(String key) throws IOException {
        Path target = resolveExisting(key);
        return Files.newInputStream(target, StandardOpenOption.READ, LinkOption.NOFOLLOW_LINKS);
    }

    @Override
    public boolean exists(String key) {
        try {
            return Files.exists(resolveExisting(key), LinkOption.NOFOLLOW_LINKS);
        } catch (IOException | IllegalArgumentException e) {
            return false;
        }
    }

    @Override
    public long size(String key) throws IOException {
        return Files.size(resolveExisting(key));
    }

    private Path resolve(String key) {
        if (key == null || key.isBlank()) {
            throw new IllegalArgumentException("archive key must not be blank");
        }
        Path target = root.resolve(key).normalize();
        if (!target.startsWith(root)) {
            throw new IllegalArgumentException("archive key escapes root: " + key);
        }
        return target;
    }

    private Path resolveExisting(String key) throws IOException {
        Path target = resolve(key);
        rejectSymlinkPath(target);
        return target;
    }

    private void createDirectoriesInsideRoot(Path directory) throws IOException {
        rejectEscapedPath(directory);
        if (Files.notExists(root, LinkOption.NOFOLLOW_LINKS)) {
            Files.createDirectories(root);
        }
        Path current = root;
        Path relative = root.relativize(directory);
        for (Path segment : relative) {
            current = current.resolve(segment);
            if (Files.isSymbolicLink(current)) {
                throw new IOException("archive path traverses symlink: " + current);
            }
            if (Files.notExists(current, LinkOption.NOFOLLOW_LINKS)) {
                Files.createDirectory(current);
            } else if (!Files.isDirectory(current, LinkOption.NOFOLLOW_LINKS)) {
                throw new IOException("archive path segment is not a directory: " + current);
            }
        }
    }

    private void rejectSymlinkPath(Path target) throws IOException {
        rejectEscapedPath(target);
        Path current = root;
        Path relative = root.relativize(target);
        for (Path segment : relative) {
            current = current.resolve(segment);
            if (Files.isSymbolicLink(current)) {
                throw new IOException("archive path traverses symlink: " + current);
            }
            if (Files.notExists(current, LinkOption.NOFOLLOW_LINKS)) {
                return;
            }
        }
    }

    private void rejectEscapedPath(Path path) {
        if (!path.normalize().startsWith(root)) {
            throw new IllegalArgumentException("archive path escapes root: " + path);
        }
    }

    private static Path rootPath(String value) {
        if (value == null || value.isBlank()) {
            return Path.of(System.getProperty("java.io.tmpdir"), "lingniu-archive");
        }
        if (value.startsWith("file://")) {
            return Path.of(URI.create(value));
        }
        return Path.of(value);
    }
}
