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.HashMap;
import java.util.Map;
import java.util.UUID;

/**
 * Codec 测试使用的内存内容存储。
 */
public class InMemorySkillContentStore implements SkillContentStore {

    private final Map<String, StoredContent> contents = new HashMap<>();
    private final Map<String, byte[]> stages = new HashMap<>();

    @Override
    public String put(byte[] bytes) {
        byte[] copy = bytes == null ? new byte[0] : Arrays.copyOf(bytes, bytes.length);
        String ref = SkillHashes.sha256Ref(copy);
        StoredContent stored = contents.get(ref);
        if (stored == null) {
            contents.put(ref, new StoredContent(copy));
        } else {
            stored.references++;
        }
        return ref;
    }

    @Override
    public SkillContentStage stage(InputStream inputStream, long maxBytes) {
        byte[] bytes = readBounded(inputStream, maxBytes);
        String stageId = UUID.randomUUID().toString();
        String hash = SkillHashes.sha256Hex(bytes);
        stages.put(stageId, bytes);
        return new SkillContentStage(stageId, "sha256:" + hash, hash, bytes.length);
    }

    @Override
    public String commit(SkillContentStage stage) {
        byte[] bytes = stage == null ? null : stages.remove(stage.getStageId());
        if (bytes == null) {
            throw new SkillException("Skill content stage does not exist.");
        }
        return put(bytes);
    }

    @Override
    public void rollback(SkillContentStage stage) {
        if (stage != null) {
            stages.remove(stage.getStageId());
        }
    }

    @Override
    public void retain(String contentRef) {
        StoredContent stored = require(contentRef);
        stored.references++;
    }

    @Override
    public void release(String contentRef) {
        StoredContent stored = contents.get(contentRef);
        if (stored != null && --stored.references <= 0) {
            contents.remove(contentRef);
        }
    }

    @Override
    public InputStream open(String contentRef) {
        StoredContent stored = require(contentRef);
        return new ByteArrayInputStream(Arrays.copyOf(stored.bytes, stored.bytes.length));
    }

    @Override
    public boolean exists(String contentRef) {
        return contents.containsKey(contentRef);
    }

    private StoredContent require(String contentRef) {
        StoredContent stored = contents.get(contentRef);
        if (stored == null) {
            throw new SkillException("Skill content does not exist: " + contentRef);
        }
        return stored;
    }

    private static byte[] readBounded(InputStream inputStream, long maxBytes) {
        try {
            ByteArrayOutputStream output = new ByteArrayOutputStream();
            byte[] buffer = new byte[8_192];
            long total = 0;
            int read;
            while ((read = inputStream.read(buffer)) >= 0) {
                total += read;
                if (total > maxBytes) {
                    throw new SkillException("Skill content exceeds limit.");
                }
                output.write(buffer, 0, read);
            }
            return output.toByteArray();
        } catch (IOException exception) {
            throw new SkillException("Failed to read Skill content.", exception);
        }
    }

    private static final class StoredContent {

        private final byte[] bytes;
        private int references = 1;

        private StoredContent(byte[] bytes) {
            this.bytes = bytes;
        }
    }
}
