package com.easyagents.skill.store.memory;

import com.easyagents.skill.exception.SkillException;
import com.easyagents.skill.store.SkillContentStage;
import com.easyagents.skill.store.SkillContentStore;
import com.easyagents.skill.util.SkillHashes;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.Arrays;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 基于内存的 Skill 内容存储，支持真实暂存与引用计数，适用于测试和轻量场景。
 */
public class InMemorySkillContentStore implements SkillContentStore {

    private final ConcurrentMap<String, StoredContent> contents = new ConcurrentHashMap<>();
    private final ConcurrentMap<String, byte[]> stagedContents = new ConcurrentHashMap<>();

    /**
     * 保存内容并持有一个引用。
     *
     * @param bytes 内容字节
     * @return 内容引用
     */
    @Override
    public String put(byte[] bytes) {
        byte[] safeBytes = bytes == null ? new byte[0] : Arrays.copyOf(bytes, bytes.length);
        String contentRef = SkillHashes.sha256Ref(safeBytes);
        contents.compute(contentRef, (key, current) -> {
            if (current == null) {
                return new StoredContent(safeBytes);
            }
            if (!Arrays.equals(current.bytes, safeBytes)) {
                throw new SkillException("SHA-256 collision detected for Skill content: " + contentRef);
            }
            current.references.incrementAndGet();
            return current;
        });
        return contentRef;
    }

    /**
     * 流式保存内容并持有一个引用。
     *
     * @param inputStream 内容流
     * @param maxBytes 最大允许字节数
     * @return 内容引用
     */
    @Override
    public String put(InputStream inputStream, long maxBytes) {
        return put(readBounded(inputStream, maxBytes));
    }

    /**
     * 将内容写入独立暂存区。
     *
     * @param inputStream 内容流
     * @param maxBytes 最大允许字节数
     * @return 暂存结果
     */
    @Override
    public SkillContentStage stage(InputStream inputStream, long maxBytes) {
        byte[] bytes = readBounded(inputStream, maxBytes);
        String hash = SkillHashes.sha256Hex(bytes);
        String stageId = UUID.randomUUID().toString();
        stagedContents.put(stageId, bytes);
        return new SkillContentStage(stageId, "sha256:" + hash, hash, bytes.length, false);
    }

    /**
     * 原子地将暂存内容转为正式引用。
     *
     * @param stage 暂存结果
     * @return 正式内容引用
     */
    @Override
    public String commit(SkillContentStage stage) {
        if (stage == null) {
            throw new SkillException("Skill content stage is required.");
        }
        byte[] bytes = stagedContents.remove(stage.getStageId());
        if (bytes == null) {
            throw new SkillException("Skill content stage does not exist: " + stage.getStageId());
        }
        String contentRef = put(bytes);
        if (!contentRef.equals(stage.getContentRef())) {
            release(contentRef);
            throw new SkillException("Skill staged content hash changed before commit.");
        }
        return contentRef;
    }

    /**
     * 删除暂存内容。
     *
     * @param stage 暂存结果
     */
    @Override
    public void rollback(SkillContentStage stage) {
        if (stage != null) {
            stagedContents.remove(stage.getStageId());
        }
    }

    /**
     * 增加正式内容引用计数。
     *
     * @param contentRef 内容引用
     */
    @Override
    public void retain(String contentRef) {
        contents.compute(contentRef, (key, content) -> {
            if (content == null) {
                throw new SkillException("Skill content does not exist: " + contentRef);
            }
            content.references.incrementAndGet();
            return content;
        });
    }

    /**
     * 释放正式内容引用并在归零时删除内容。
     *
     * @param contentRef 内容引用
     */
    @Override
    public void release(String contentRef) {
        contents.computeIfPresent(contentRef, (key, content) ->
                content.references.decrementAndGet() <= 0 ? null : content);
    }

    /**
     * 打开内容流。
     *
     * @param contentRef 内容引用
     * @return 内容流
     */
    @Override
    public InputStream open(String contentRef) {
        return new ByteArrayInputStream(readAllBytes(contentRef));
    }

    /**
     * 读取全部内容。
     *
     * @param contentRef 内容引用
     * @return 内容副本
     */
    @Override
    public byte[] readAllBytes(String contentRef) {
        StoredContent content = contents.get(contentRef);
        if (content == null) {
            throw new SkillException("Skill content does not exist: " + contentRef);
        }
        return Arrays.copyOf(content.bytes, content.bytes.length);
    }

    /**
     * 判断内容是否存在。
     *
     * @param contentRef 内容引用
     * @return 存在时为 true
     */
    @Override
    public boolean exists(String contentRef) {
        return contents.containsKey(contentRef);
    }

    private static byte[] readBounded(InputStream inputStream, long maxBytes) {
        if (inputStream == null || maxBytes < 0) {
            throw new SkillException("Valid Skill content stream and limit are required.");
        }
        try {
            ByteArrayOutputStream output = new ByteArrayOutputStream((int) Math.min(maxBytes, 8_192));
            byte[] buffer = new byte[8_192];
            long total = 0;
            int read;
            while ((read = inputStream.read(buffer)) >= 0) {
                if (read == 0) {
                    continue;
                }
                total += read;
                if (total > maxBytes) {
                    throw new SkillException("Skill content exceeds " + maxBytes + " bytes.");
                }
                output.write(buffer, 0, read);
            }
            return output.toByteArray();
        } catch (IOException e) {
            throw new SkillException("Failed to read Skill content stream.", e);
        }
    }

    private static final class StoredContent {

        private final byte[] bytes;
        private final AtomicInteger references = new AtomicInteger(1);

        private StoredContent(byte[] bytes) {
            this.bytes = Arrays.copyOf(bytes, bytes.length);
        }
    }
}
