package com.gzzm.lobster.sandbox;

import com.gzzm.lobster.config.LobsterConfig;
import com.gzzm.platform.commons.Tools;
import net.cyan.arachne.annotation.Service;
import net.cyan.nest.annotation.Inject;

import java.io.IOException;
import java.nio.file.FileSystems;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.nio.file.attribute.PosixFilePermission;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.EnumSet;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.stream.Stream;

/**
 * Fixed-slot sandbox pool.
 *
 * <p>Each slot owns stable host directories and at most one pre-created Docker
 * container. A slot is leased by one code_exec at a time. After the caller has
 * harvested outputs, release removes the used container, clears the directories,
 * and asynchronously pre-creates a fresh container for the same execution
 * profile. The pool never runs multiple user tasks inside one running container.
 */
@Service
public class SandboxPoolService {

    @Inject private DockerRunner dockerRunner;

    /**
     * nest DI creates separate shell instances for injected services. Pool state
     * must therefore be JVM-shared, otherwise SandboxService can initialize one
     * pool while the admin API reads an empty pool from another instance.
     */
    private static final Object initLock = new Object();
    private static volatile boolean initialized;
    private static volatile BlockingQueue<Slot> available;
    private static volatile List<Slot> slots = Collections.emptyList();

    private static final ExecutorService prewarmPool = Executors.newCachedThreadPool(r -> {
        Thread t = new Thread(r, "sandbox-pool-prewarm");
        t.setDaemon(true);
        return t;
    });

    public void warmUp() {
        if (!LobsterConfig.isSandboxPoolEnabled()) return;
        ensureInitialized();
    }

    public Lease acquire() throws SandboxException {
        ensureInitialized();
        try {
            long deadline = System.currentTimeMillis() + TimeUnit.SECONDS.toMillis(30);
            while (true) {
                long remaining = deadline - System.currentTimeMillis();
                if (remaining <= 0) {
                    throw new SandboxException("sandbox.pool_busy",
                            "sandbox pool has no free slot within 30s");
                }
                Slot slot = available.poll(remaining, TimeUnit.MILLISECONDS);
                if (slot == null) {
                    throw new SandboxException("sandbox.pool_busy",
                            "sandbox pool has no free slot within 30s");
                }
                if (!slot.leased.compareAndSet(false, true)) {
                    continue;
                }
                try {
                    synchronized (slot) {
                        // A pre-created container already has bind mounts attached to the current
                        // slot directories. Deleting and recreating those directories here leaves
                        // the container mounted to the old, now-unlinked paths, so entry.js/entry.py
                        // written by SandboxService becomes invisible inside /work.
                        if (slot.containerId == null) {
                            resetSlotDirs(slot);
                        }
                    }
                } catch (IOException e) {
                    slot.leased.set(false);
                    offer(slot);
                    throw new SandboxException("sandbox.pool_reset",
                            "failed to reset sandbox pool slot " + slot.root + ": " + e.getMessage(), e);
                }
                return new Lease(slot);
            }
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new SandboxException("sandbox.pool_busy", "interrupted while waiting for sandbox slot", e);
        }
    }

    public void shutdown() {
        prewarmPool.shutdownNow();
        for (Slot slot : slots) {
            synchronized (slot) {
                dockerRunner.removeContainerQuiet(slot.containerId == null ? slot.containerName : slot.containerId);
                slot.containerId = null;
                slot.signature = null;
                slot.leased.set(false);
            }
        }
        BlockingQueue<Slot> q = available;
        if (q != null) q.clear();
    }

    public Map<String, Object> snapshot() {
        Map<String, Object> out = new LinkedHashMap<>();
        out.put("enabled", Boolean.valueOf(LobsterConfig.isSandboxPoolEnabled()));
        out.put("initialized", Boolean.valueOf(initialized));
        out.put("size", Integer.valueOf(Math.max(1, LobsterConfig.getSandboxPoolSize())));
        out.put("root", LobsterConfig.getSandboxPoolRoot());
        BlockingQueue<Slot> q = available;
        out.put("availableQueueSize", Integer.valueOf(q == null ? 0 : q.size()));
        List<Map<String, Object>> rows = new ArrayList<>();
        for (Slot slot : slots) {
            Map<String, Object> row = new LinkedHashMap<>();
            synchronized (slot) {
                row.put("containerName", slot.containerName);
                row.put("root", slot.root.toString());
                row.put("inputs", slot.inputs.toString());
                row.put("outputs", slot.outputs.toString());
                row.put("work", slot.work.toString());
                row.put("leased", Boolean.valueOf(slot.leased.get()));
                row.put("containerId", slot.containerId);
                row.put("signature", slot.signature);
            }
            row.put("dockerStatus", "");
            rows.add(row);
        }
        out.put("slots", rows);
        return out;
    }

    public DockerRunner.RunOutcome start(Lease lease, DockerRunner.RunSpec spec) throws SandboxException {
        if (lease == null || lease.slot == null) {
            throw new SandboxException("sandbox.pool", "missing sandbox pool lease");
        }
        Slot slot = lease.slot;
        String containerId = prepareContainer(slot, spec);
        if (spec.onContainerCreated != null) {
            try { spec.onContainerCreated.accept(containerId); }
            catch (Throwable ignore) { /* do not block start */ }
        }
        if (spec.cancelCheck != null) {
            try {
                if (spec.cancelCheck.getAsBoolean()) {
                    DockerRunner.RunOutcome out = new DockerRunner.RunOutcome();
                    out.cancelled = true;
                    out.exitCode = DockerRunner.EXIT_TIMEOUT;
                    return out;
                }
            } catch (Throwable ignore) { /* ignore */ }
        }
        return dockerRunner.startCreated(spec, containerId);
    }

    public void release(Lease lease, DockerRunner.RunSpec lastSpec) {
        if (lease == null || lease.slot == null) return;
        if (!lease.released.compareAndSet(false, true)) return;
        Slot slot = lease.slot;
        String oldContainer;
        synchronized (slot) {
            oldContainer = slot.containerId;
            slot.containerId = null;
        }
        dockerRunner.removeContainerQuiet(oldContainer == null ? slot.containerName : oldContainer);
        try { resetSlotDirs(slot); } catch (Throwable t) {
            try { Tools.log("[SandboxPool] reset failed for " + slot.root, t); } catch (Throwable ignore) { /* ignore */ }
        }
        if (lastSpec == null) {
            slot.leased.set(false);
            offer(slot);
            return;
        }
        DockerRunner.RunSpec prewarmSpec = copyForCreate(lastSpec, slot.containerName);
        String signature = signature(prewarmSpec);
        try {
            prewarmPool.submit(() -> {
                try {
                    synchronized (slot) {
                        slot.containerId = createFresh(slot, prewarmSpec);
                        slot.signature = signature;
                    }
                } catch (Throwable t) {
                    synchronized (slot) {
                        slot.containerId = null;
                        slot.signature = null;
                    }
                    try { Tools.log("[SandboxPool] prewarm failed for " + slot.containerName, t); }
                    catch (Throwable ignore) { /* ignore */ }
                } finally {
                    slot.leased.set(false);
                    offer(slot);
                }
            });
        } catch (RejectedExecutionException e) {
            slot.leased.set(false);
            offer(slot);
        }
    }

    public Path inputs(Lease lease) { return lease.slot.inputs; }
    public Path outputs(Lease lease) { return lease.slot.outputs; }
    public Path work(Lease lease) { return lease.slot.work; }
    public Path root(Lease lease) { return lease.slot.root; }
    public String containerName(Lease lease) { return lease.slot.containerName; }

    public static final class Lease {
        private final Slot slot;
        private final AtomicBoolean released = new AtomicBoolean(false);
        private Lease(Slot slot) { this.slot = slot; }
    }

    private static final class Slot {
        final String containerName;
        final Path root;
        final Path inputs;
        final Path outputs;
        final Path work;
        final AtomicBoolean leased = new AtomicBoolean(false);
        String containerId;
        String signature;

        Slot(int index, Path root) {
            this.containerName = String.format("lobster-sbx-pool-%03d", index);
            this.root = root;
            this.inputs = root.resolve("inputs");
            this.outputs = root.resolve("outputs");
            this.work = root.resolve("work");
        }
    }

    private void ensureInitialized() {
        if (initialized) return;
        synchronized (initLock) {
            if (initialized) return;
            int size = Math.max(1, LobsterConfig.getSandboxPoolSize());
            BlockingQueue<Slot> q = new LinkedBlockingQueue<>();
            List<Slot> initializedSlots = new ArrayList<>();
            Path root = Paths.get(LobsterConfig.getSandboxPoolRoot()).toAbsolutePath();
            for (int i = 0; i < size; i++) {
                Slot slot = new Slot(i, root.resolve(String.format("slot-%03d", i)));
                try {
                    resetSlotDirs(slot);
                } catch (IOException e) {
                    throw new SandboxException("sandbox.pool_init",
                            "failed to initialize sandbox pool slot " + slot.root + ": " + e.getMessage(), e);
                }
                initializedSlots.add(slot);
                q.offer(slot);
            }
            available = q;
            slots = Collections.unmodifiableList(initializedSlots);
            initialized = true;
            for (Slot slot : initializedSlots) {
                scheduleInitialPrewarm(slot);
            }
        }
    }

    private String prepareContainer(Slot slot, DockerRunner.RunSpec spec) throws SandboxException {
        DockerRunner.RunSpec createSpec = copyForCreate(spec, slot.containerName);
        String sig = signature(createSpec);
        synchronized (slot) {
            if (slot.containerId != null && sig.equals(slot.signature)) {
                if (dockerRunner.isReusableCreatedContainer(slot.containerId)) {
                    return slot.containerId;
                }
                try { Tools.log("[SandboxPool] prewarmed container missing/stale; recreating " + slot.containerName); }
                catch (Throwable ignore) { /* ignore */ }
            }
            dockerRunner.removeContainerQuiet(slot.containerId == null ? slot.containerName : slot.containerId);
            slot.containerId = createFresh(slot, createSpec);
            slot.signature = sig;
            return slot.containerId;
        }
    }

    private String createFresh(Slot slot, DockerRunner.RunSpec createSpec) throws SandboxException {
        dockerRunner.removeContainerQuiet(slot.containerName);
        createSpec.containerName = slot.containerName;
        return dockerRunner.create(createSpec);
    }

    private void scheduleInitialPrewarm(Slot slot) {
        DockerRunner.RunSpec prewarmSpec = defaultPrewarmSpec(slot);
        String sig = signature(prewarmSpec);
        prewarmPool.submit(() -> {
            try {
                synchronized (slot) {
                    if (slot.leased.get()) {
                        return;
                    }
                    slot.containerId = createFresh(slot, prewarmSpec);
                    slot.signature = sig;
                }
                try { Tools.log("[SandboxPool] prewarmed " + slot.containerName); }
                catch (Throwable ignore) { /* ignore */ }
            } catch (Throwable t) {
                synchronized (slot) {
                    slot.containerId = null;
                    slot.signature = null;
                }
                try { Tools.log("[SandboxPool] initial prewarm failed for " + slot.containerName, t); }
                catch (Throwable ignore) { /* ignore */ }
            }
        });
    }

    private static DockerRunner.RunSpec defaultPrewarmSpec(Slot slot) {
        DockerRunner.RunSpec spec = new DockerRunner.RunSpec();
        spec.image = LobsterConfig.getSandboxImage();
        spec.containerName = slot.containerName;
        spec.user = LobsterConfig.getSandboxUid() + ":" + LobsterConfig.getSandboxUid();
        spec.memoryMb = LobsterConfig.getSandboxMemoryMb();
        spec.cpus = LobsterConfig.getSandboxCpus();
        spec.pidsLimit = LobsterConfig.getSandboxPidsLimit();
        spec.walltimeSec = LobsterConfig.getSandboxDefaultTimeoutSec();
        spec.mounts.add(new DockerRunner.Mount(slot.inputs.toString(), "/inputs", true));
        spec.mounts.add(new DockerRunner.Mount(slot.outputs.toString(), "/outputs", false));
        spec.mounts.add(new DockerRunner.Mount(slot.work.toString(), "/work", false));
        spec.tmpfs.add("/tmp:size=128m");
        int uid = LobsterConfig.getSandboxUid();
        spec.tmpfs.add("/home/sandbox:size=64m,uid=" + uid + ",gid=" + uid + ",mode=0700");
        spec.args.add("/work/entry.py");
        return spec;
    }

    private void offer(Slot slot) {
        try { available.offer(slot); } catch (Throwable ignore) { /* impossible for unbounded queue */ }
    }

    private static DockerRunner.RunSpec copyForCreate(DockerRunner.RunSpec spec, String containerName) {
        DockerRunner.RunSpec copy = new DockerRunner.RunSpec();
        copy.image = spec.image;
        copy.containerName = containerName;
        copy.user = spec.user;
        copy.memoryMb = spec.memoryMb;
        copy.cpus = spec.cpus;
        copy.pidsLimit = spec.pidsLimit;
        copy.walltimeSec = spec.walltimeSec;
        copy.entrypoint = spec.entrypoint;
        copy.args.addAll(spec.args);
        copy.tmpfs.addAll(spec.tmpfs);
        for (DockerRunner.Mount m : spec.mounts) {
            copy.mounts.add(new DockerRunner.Mount(m.host, m.container, m.readOnly));
        }
        return copy;
    }

    private static String signature(DockerRunner.RunSpec spec) {
        StringBuilder sb = new StringBuilder();
        sb.append("image=").append(spec.image).append('\n');
        sb.append("user=").append(spec.user).append('\n');
        sb.append("memory=").append(spec.memoryMb).append('\n');
        sb.append("cpus=").append(spec.cpus).append('\n');
        sb.append("pids=").append(spec.pidsLimit).append('\n');
        sb.append("entrypoint=").append(spec.entrypoint).append('\n');
        sb.append("args=").append(spec.args).append('\n');
        sb.append("tmpfs=").append(spec.tmpfs).append('\n');
        for (DockerRunner.Mount m : spec.mounts) {
            sb.append("mount=").append(m.host).append(':')
                    .append(m.container).append(':').append(m.readOnly).append('\n');
        }
        return sb.toString();
    }

    private static void resetSlotDirs(Slot slot) throws IOException {
        Files.createDirectories(slot.inputs);
        Files.createDirectories(slot.outputs);
        Files.createDirectories(slot.work);
        deleteChildren(slot.inputs);
        deleteChildren(slot.outputs);
        deleteChildren(slot.work);
        relaxPermissions(slot.inputs);
        relaxPermissions(slot.outputs);
        relaxPermissions(slot.work);
    }

    private static void deleteChildren(Path root) throws IOException {
        if (!Files.exists(root)) return;
        try (Stream<Path> w = Files.walk(root)) {
            w.sorted(Comparator.reverseOrder()).forEach(p -> {
                if (p.equals(root)) return;
                try { Files.deleteIfExists(p); } catch (IOException ignore) { /* best effort */ }
            });
        }
    }

    private static void relaxPermissions(Path dir) {
        if (!FileSystems.getDefault().supportedFileAttributeViews().contains("posix")) return;
        try {
            Set<PosixFilePermission> perms = EnumSet.of(
                    PosixFilePermission.OWNER_READ, PosixFilePermission.OWNER_WRITE, PosixFilePermission.OWNER_EXECUTE,
                    PosixFilePermission.GROUP_READ, PosixFilePermission.GROUP_WRITE, PosixFilePermission.GROUP_EXECUTE,
                    PosixFilePermission.OTHERS_READ, PosixFilePermission.OTHERS_WRITE, PosixFilePermission.OTHERS_EXECUTE);
            Files.setPosixFilePermissions(dir, perms);
        } catch (Throwable ignore) { /* best effort */ }
    }
}