package com.easyagents.agent.runtime.tool.operate;

import com.easyagents.agent.runtime.AgentRuntimeException;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.LinkOption;
import java.nio.file.Path;
import java.util.Map;
import java.util.HashSet;
import java.util.Set;
import java.util.stream.Stream;

/**
 * 工作区容量与文件数量校验器。
 */
final class WorkspaceQuotaGuard {

    private static final long MAX_SCANNED_ENTRIES = 100_000L;
    private static final int DEFAULT_ARCHIVE_ENTRY_LIMIT = 10_000;
    private static final long DEFAULT_ARCHIVE_TOTAL_LIMIT = 512L * 1024L * 1024L;
    private static final long DEFAULT_ARCHIVE_FILE_LIMIT = 64L * 1024L * 1024L;

    private final WorkspacePathGuard pathGuard;
    private final WorkspaceQuotaLimits limits;
    private final WorkspaceQuotaHook hook;

    /**
     * 创建配额校验器。
     *
     * @param pathGuard 路径保护器
     * @param limits 配额限制
     * @param hook 业务侧附加校验 Hook
     */
    WorkspaceQuotaGuard(WorkspacePathGuard pathGuard,
                        WorkspaceQuotaLimits limits,
                        WorkspaceQuotaHook hook) {
        this.pathGuard = pathGuard;
        this.limits = limits == null ? WorkspaceQuotaLimits.unlimited() : limits;
        this.hook = hook == null ? WorkspaceQuotaHook.noop() : hook;
    }

    /**
     * 校验文件是否允许被完整读取。
     *
     * @param target 目标普通文件
     */
    void validateFullRead(Path target) {
        try {
            long size = Files.size(target);
            if (limits.getMaxReadSize() > 0 && size > limits.getMaxReadSize()) {
                throw new WorkspaceToolException("WORKSPACE_QUOTA_EXCEEDED",
                        "Workspace full-file read exceeds max-read-size.", false);
            }
            hook.beforeRead(pathGuard.root(), target, size);
        } catch (IOException error) {
            throw new WorkspaceToolException("WORKSPACE_IO_FAILED",
                    "Workspace file size cannot be inspected.", true, error);
        }
    }

    /**
     * 记录一次范围读取并调用业务侧配额 Hook。
     *
     * @param target 目标文件
     * @param readBytes 实际返回字节数
     */
    void validateRangeRead(Path target, long readBytes) {
        if (limits.getMaxReadSize() > 0 && readBytes > limits.getMaxReadSize()) {
            throw new WorkspaceToolException("WORKSPACE_QUOTA_EXCEEDED",
                    "Workspace range read exceeds max-read-size.", false);
        }
        hook.beforeRead(pathGuard.root(), target, readBytes);
    }

    /**
     * 获取范围读取字节上限。
     *
     * @return 字节上限，零表示使用 Runtime 固定安全上限
     */
    long maxReadSize() {
        return limits.getMaxReadSize() > 0 ? limits.getMaxReadSize() : 2L * 1024L * 1024L;
    }

    /**
     * 获取单层目录最大返回条目数。
     *
     * @return 最大条目数
     */
    int maxDirectoryEntries() {
        long configured = limits.getMaxFileCount();
        return configured > 0 ? (int) Math.min(configured, 1000) : 1000;
    }

    /**
     * 获取安全归档单次最大条目数。
     *
     * @return 条目数上限
     */
    int maxArchiveEntries() {
        long configured = limits.getMaxFileCount();
        return configured > 0
                ? (int) Math.min(configured, DEFAULT_ARCHIVE_ENTRY_LIMIT)
                : DEFAULT_ARCHIVE_ENTRY_LIMIT;
    }

    /**
     * 获取安全归档展开总量上限。
     *
     * @return 展开总字节数上限
     */
    long maxArchiveTotalSize() {
        long configured = limits.getMaxTotalSize();
        return configured > 0 ? Math.min(configured, DEFAULT_ARCHIVE_TOTAL_LIMIT) : DEFAULT_ARCHIVE_TOTAL_LIMIT;
    }

    /**
     * 获取安全归档单文件上限。
     *
     * @return 单文件字节数上限
     */
    long maxArchiveSingleFileSize() {
        long configured = limits.getMaxSingleFileSize();
        return configured > 0 ? Math.min(configured, DEFAULT_ARCHIVE_FILE_LIMIT) : DEFAULT_ARCHIVE_FILE_LIMIT;
    }

    /**
     * 校验单个文件变更后的工作区配额。
     *
     * @param target 目标文件
     * @param resultingBytes 变更后的文件字节数，删除时为零
     */
    void validateWrite(Path target, long resultingBytes) {
        validateBatch(Map.of(target, resultingBytes));
    }

    /**
     * 校验一批文件变更后的工作区配额。
     *
     * @param resultingSizes 目标路径到变更后字节数的映射，负数表示删除
     */
    void validateBatch(Map<Path, Long> resultingSizes) {
        if (resultingSizes == null || resultingSizes.isEmpty()) {
            return;
        }
        WorkspaceUsage usage = scanUsage();
        long projectedSize = usage.totalSize();
        long projectedCount = usage.entryCount();
        Set<Path> plannedEntries = new HashSet<>();
        for (Map.Entry<Path, Long> entry : resultingSizes.entrySet()) {
            Path target = entry.getKey();
            long resultingBytes = entry.getValue() == null ? 0 : entry.getValue();
            boolean exists = Files.isRegularFile(target, LinkOption.NOFOLLOW_LINKS);
            long previousBytes = sizeIfRegular(target);
            projectedSize -= previousBytes;
            if (resultingBytes < 0) {
                if (exists) {
                    projectedCount--;
                }
                hook.beforeWrite(pathGuard.root(), target, previousBytes, 0);
                continue;
            }
            if (limits.getMaxSingleFileSize() > 0 && resultingBytes > limits.getMaxSingleFileSize()) {
                throw new WorkspaceToolException("WORKSPACE_QUOTA_EXCEEDED",
                        "Workspace file exceeds max-single-file-size.", false);
            }
            try {
                projectedSize = Math.addExact(projectedSize, resultingBytes);
            } catch (ArithmeticException error) {
                throw new WorkspaceToolException("WORKSPACE_QUOTA_EXCEEDED",
                        "Workspace exceeds max-total-size.", false, error);
            }
            if (!exists) {
                if (plannedEntries.add(target)) {
                    projectedCount++;
                }
                Path parent = target.getParent();
                while (parent != null && !parent.equals(pathGuard.root())) {
                    if (!Files.exists(parent, LinkOption.NOFOLLOW_LINKS) && plannedEntries.add(parent)) {
                        projectedCount++;
                    }
                    parent = parent.getParent();
                }
            }
            hook.beforeWrite(pathGuard.root(), target, previousBytes, resultingBytes);
        }
        if (limits.getMaxTotalSize() > 0 && projectedSize > limits.getMaxTotalSize()) {
            throw new WorkspaceToolException("WORKSPACE_QUOTA_EXCEEDED",
                    "Workspace exceeds max-total-size.", false);
        }
        if (limits.getMaxFileCount() > 0 && projectedCount > limits.getMaxFileCount()) {
            throw new WorkspaceToolException("WORKSPACE_QUOTA_EXCEEDED",
                    "Workspace exceeds max-file-count.", false);
        }
    }

    /**
     * 校验当前工作区已处于配额范围内。
     */
    void validateCurrentUsage() {
        WorkspaceUsage usage = scanUsage();
        if (limits.getMaxTotalSize() > 0 && usage.totalSize() > limits.getMaxTotalSize()) {
            throw new WorkspaceToolException("WORKSPACE_QUOTA_EXCEEDED",
                    "Workspace exceeds max-total-size.", false);
        }
        if (limits.getMaxFileCount() > 0 && usage.entryCount() > limits.getMaxFileCount()) {
            throw new WorkspaceToolException("WORKSPACE_QUOTA_EXCEEDED",
                    "Workspace exceeds max-file-count.", false);
        }
    }

    private WorkspaceUsage scanUsage() {
        long totalSize = 0;
        long entryCount = 0;
        try (Stream<Path> paths = Files.walk(pathGuard.root())) {
            for (Path path : (Iterable<Path>) paths::iterator) {
                if (path.equals(pathGuard.root())) {
                    continue;
                }
                entryCount++;
                if (entryCount > MAX_SCANNED_ENTRIES) {
                    throw new WorkspaceToolException("WORKSPACE_QUOTA_EXCEEDED",
                            "Workspace contains too many entries to inspect safely.", false);
                }
                if (Files.isSymbolicLink(path)) {
                    throw new WorkspaceToolException("WORKSPACE_PATH_INVALID",
                            "Workspace contains a symbolic link.", false);
                }
                if (Files.isRegularFile(path, LinkOption.NOFOLLOW_LINKS)) {
                    totalSize = Math.addExact(totalSize, Files.size(path));
                } else if (!Files.isDirectory(path, LinkOption.NOFOLLOW_LINKS)) {
                    throw new WorkspaceToolException("FILE_TYPE_INVALID",
                            "Workspace contains a non-regular entry.", false);
                }
            }
            return new WorkspaceUsage(totalSize, entryCount);
        } catch (IOException | ArithmeticException error) {
            throw new WorkspaceToolException("WORKSPACE_IO_FAILED",
                    "Workspace usage cannot be inspected.", true, error);
        }
    }

    private long sizeIfRegular(Path target) {
        if (!Files.isRegularFile(target, LinkOption.NOFOLLOW_LINKS)) {
            return 0;
        }
        try {
            return Files.size(target);
        } catch (IOException error) {
            throw new WorkspaceToolException("WORKSPACE_IO_FAILED",
                    "Workspace file size cannot be inspected.", true, error);
        }
    }

    /**
     * 工作区当前使用量。
     *
     * @param totalSize 普通文件总字节数
     * @param entryCount 文件与目录条目数量，不包含工作区根
     */
    private record WorkspaceUsage(long totalSize, long entryCount) {
    }
}
