#include "container.hpp"

#include <fcntl.h>
#include <sys/stat.h>
#include <unistd.h>

#include <algorithm>
#include <cstdio>
#include <cctype>
#include <ctime>
#include <map>
#include <sstream>

#include "spec.hpp"
#include "util.hpp"

namespace fs = std::filesystem;

namespace doos {

namespace {

std::string slug(std::string_view hostDir) {
    std::string s(hostDir);
    if (!s.empty() && s.front() == '/') s.erase(0, 1);
    return s;
}

std::map<std::string, std::string> parseKv(const std::string& text) {
    std::map<std::string, std::string> kv;
    std::istringstream in(text);
    for (std::string line; std::getline(in, line);) {
        auto eq = line.find('=');
        if (eq != std::string::npos) kv[line.substr(0, eq)] = line.substr(eq + 1);
    }
    return kv;
}

// starttime (field 22 of /proc/PID/stat), used to defeat pid reuse.
std::optional<std::string> startTime(pid_t pid) {
    try {
        std::string st = fsutil::readFile("/proc/" + std::to_string(pid) + "/stat");
        auto rp = st.rfind(')');
        if (rp == std::string::npos) return std::nullopt;
        std::istringstream in(st.substr(rp + 2));
        std::string tok;
        for (int i = 3; i <= 22; ++i) {
            if (!(in >> tok)) return std::nullopt;
        }
        return tok;
    } catch (const std::exception&) {
        return std::nullopt;
    }
}

}  // namespace

void Container::validateName(const std::string& name) {
    bool ok = !name.empty() && name.size() <= spec::kMaxNameLen &&
              (std::islower(static_cast<unsigned char>(name[0])) || std::isdigit(static_cast<unsigned char>(name[0])));
    for (char c : name) {
        ok = ok && (std::islower(static_cast<unsigned char>(c)) || std::isdigit(static_cast<unsigned char>(c)) ||
                    c == '-' || c == '_');
    }
    if (!ok) {
        throw std::runtime_error("invalid container name '" + name + "': use [a-z0-9][a-z0-9_-]{0,31}");
    }
}

fs::path Container::layer(const std::string& hostDir) const { return dir_ / "layers" / slug(hostDir); }
fs::path Container::privateDir(const std::string& hostDir) const { return dir_ / "private" / slug(hostDir); }

fs::path Container::etcDir() const {
    return imageMode() ? layer("/root") / "upper" / "etc" : privateDir("/etc");
}

void Container::seed(const std::string& marker, const std::string& label,
                     const std::function<fsutil::SeedStats()>& work) const {
    fs::path m = dir_ / ("seeded." + marker);
    if (fs::exists(m)) return;
    fsutil::SeedStats st = work();
    std::fprintf(stderr, "doos: seeded %s: %zu entries%s%s\n", label.c_str(), st.made,
                 st.skipped ? ", unreadable skipped: " : "", st.skipped ? std::to_string(st.skipped).c_str() : "");
    fsutil::writeFile(m, "ok\n");
}

namespace {
// Names that may hold credentials are never copied, even inside allowed directories.
bool etcDenied(const fs::path& p) {
    std::string n = p.filename().string();
    if (n.size() > 4 && n.compare(n.size() - 4, 4, ".key") == 0) return true;
    for (auto d : spec::kEtcDenyNames) {
        std::string pat(d);
        if (!pat.empty() && pat.back() == '*') {
            if (n.rfind(pat.substr(0, pat.size() - 1), 0) == 0) return true;
        } else if (n == pat) {
            return true;
        }
    }
    return false;
}
}  // namespace

void Container::prepareLayout() const {
    fsutil::mkdirs(rootfs());
    if (imageMode()) {
        fsutil::mkdirs(layer("/root") / "upper");
        fsutil::mkdirs(layer("/root") / "work");
        return;
    }
    std::vector<std::string> overlays;
    auto addOverlay = [&](std::string_view d) {
        std::string dir(d);
        fsutil::mkdirs(layer(dir) / "upper");
        fsutil::mkdirs(layer(dir) / "work");
        overlays.push_back(dir);
    };
    for (auto d : spec::kOverlayDirs) {
        fsutil::mkdirs(rootfs() / slug(d));
        addOverlay(d);
    }
    for (auto d : spec::kCompatDirs) {
        std::error_code ec;
        if (fs::is_directory(fs::symlink_status(std::string(d), ec)) && !ec) {
            fsutil::mkdirs(rootfs() / slug(d));
            addOverlay(d);
        }
    }
    for (auto d : spec::kPrivateDirs) {
        fsutil::mkdirs(rootfs() / slug(d));
        fsutil::mkdirs(privateDir(std::string(d)), 0755);
    }
    for (auto d : spec::kVarDirs) fsutil::mkdirs(privateDir("/var") / d.path, d.mode);
    for (auto d : spec::kVarOverlayDirs) {
        std::error_code ec;
        if (!fs::is_directory(std::string(d), ec)) continue;
        fsutil::mkdirs(privateDir("/var") / std::string(d).substr(5));  // below /var
        addOverlay(d);
    }
    for (auto d : spec::kEmptyDirs) fsutil::mkdirs(rootfs() / slug(d));

    // Host dirs are owned by an unmapped uid (nobody) in the container: without a directory skeleton in the
    // upper layer nothing could ever be created in them (package installs, tool caches).
    for (const auto& o : overlays) {
        std::error_code ec;
        if (!fs::is_directory(o, ec)) continue;
        std::string key = slug(o);
        std::replace(key.begin(), key.end(), '/', '_');
        seed("skeleton-" + key, "directory skeleton " + o, [&] { return fsutil::mirrorDirs(o, layer(o) / "upper"); });
    }
    for (auto d : spec::kUpperCopyDirs) {
        for (const auto& o : overlays) {
            if (std::string(d) != o && std::string(d).rfind(o + "/", 0) != 0) continue;
            std::string rel = std::string(d).substr(o.size());
            std::string key = slug(d);
            std::replace(key.begin(), key.end(), '/', '_');
            fs::path target = layer(o) / "upper" / (rel.empty() ? "" : rel.substr(1));
            std::error_code ec;
            if (!fs::is_directory(std::string(d), ec) || !fs::is_directory(target, ec)) continue;
            seed("files-" + key, "files of " + std::string(d), [&] { return fsutil::copyFiles(std::string(d), target); });
        }
    }
    for (auto f : spec::kUpperLockFiles) {
        for (const auto& o : overlays) {
            if (std::string(f).rfind(o + "/", 0) != 0) continue;
            fs::path target = layer(o) / "upper" / std::string(f).substr(o.size() + 1);
            std::error_code ec;
            if (!fs::is_directory(target.parent_path(), ec) || fs::exists(target, ec)) continue;
            fsutil::writeFile(target, "");
            ::chmod(target.c_str(), 0640);
        }
    }
    for (auto d : spec::kCopiedDirs) {
        std::error_code ec;
        if (!fs::is_directory(std::string(d), ec)) continue;
        std::string key = slug(d);
        std::replace(key.begin(), key.end(), '/', '_');
        if (relaxed_) {
            seed("copy-" + key, "copy of " + std::string(d) + " (relaxed)",
                 [&] { return fsutil::copyTree(std::string(d), privateDir(std::string(d))); });
        } else {
            seed("copy-" + key, "allowlisted " + std::string(d), [&] {
                std::vector<std::string> names(spec::kEtcAllow.begin(), spec::kEtcAllow.end());
                return fsutil::copyEntries(std::string(d), privateDir(std::string(d)), names, etcDenied);
            });
        }
    }
}

namespace {
void findUppers(const fs::path& dir, std::size_t depth, std::vector<fs::path>& out) {
    std::error_code ec;
    if (depth > 6) return;
    for (const auto& e : fs::directory_iterator(dir, ec)) {
        if (!e.is_directory(ec) || e.is_symlink(ec)) continue;
        std::string n = e.path().filename().string();
        if (n == "upper") out.push_back(e.path());
        else if (n != "work") findUppers(e.path(), depth + 1, out);
    }
}
}  // namespace

std::uint64_t Container::usedBytes(std::size_t& skipped) const {
    std::vector<fs::path> dirs;
    std::error_code ec;
    if (fs::exists(dir_ / "layers", ec)) findUppers(dir_ / "layers", 0, dirs);
    if (fs::exists(dir_ / "private", ec)) dirs.push_back(dir_ / "private");
    auto scan = [&](std::size_t& sk) {
        std::uint64_t total = 0;
        for (const auto& d : dirs) total += fsutil::diskUsage(d, sk);
        return total;
    };
    if (ids_ == IdMode::Single) return scan(skipped);
    // Files of other container ids (host uids 100000+) and their mode 700 directories are invisible to the host
    // user; measuring as root of the id map sees them.
    int p[2];
    if (::pipe2(p, O_CLOEXEC) != 0) throw SysError("pipe");
    Fd rd(p[0]), wr(p[1]);
    int rc = UserNsJob::run(ids_, [&] {
        std::size_t sk = 0;
        std::uint64_t t = scan(sk);
        ::dprintf(wr.get(), "%llu %zu\n", static_cast<unsigned long long>(t), sk);
        return 0;
    });
    wr.reset();
    if (rc != 0) throw std::runtime_error("size scan failed");
    char buf[64] = {0};
    if (::read(rd.get(), buf, sizeof buf - 1) <= 0) throw std::runtime_error("size scan returned nothing");
    unsigned long long t = 0;
    std::size_t sk = 0;
    std::sscanf(buf, "%llu %zu", &t, &sk);
    skipped += sk;
    return t;
}

void Container::writeMeta() const {
    fsutil::writeFile(metaFile(), "net=" + Network::name(net_) + "\ncreated=" + std::to_string(created_) +
                                      "\nversion=" DOOS_VERSION "\nimage=" + image_ + "\nids=" + IdMap::name(ids_) +
                                      "\nsize=" + std::to_string(size_) + "\nmemory=" + std::to_string(memory_) +
                                      "\npids=" + std::to_string(pids_) + "\ncpus=" + std::to_string(milliCpus_) +
                                      "\nhardening=" + (relaxed_ ? "relaxed" : "default") + "\n");
}

void Container::setSizeCap(std::uint64_t bytes) {
    if (bytes < spec::kMinSizeBytes || bytes > spec::kMaxSizeBytes) {
        throw std::runtime_error("size must be between " + formatSize(spec::kMinSizeBytes) + " and " +
                                 formatSize(spec::kMaxSizeBytes));
    }
    size_ = bytes;
    writeMeta();
}

void Container::setLimits(std::optional<std::uint64_t> memory, std::optional<std::uint32_t> pids,
                          std::optional<std::uint32_t> milliCpus) {
    if (memory) {
        if (*memory < spec::kMinMemoryBytes || *memory > spec::kMaxMemoryBytes) {
            throw std::runtime_error("memory must be between " + formatSize(spec::kMinMemoryBytes) + " and " +
                                     formatSize(spec::kMaxMemoryBytes));
        }
        memory_ = *memory;
    }
    if (pids) {
        if (*pids < spec::kMinPids || *pids > spec::kMaxPids) {
            throw std::runtime_error("pids must be between " + std::to_string(spec::kMinPids) + " and " +
                                     std::to_string(spec::kMaxPids));
        }
        pids_ = *pids;
    }
    if (milliCpus) {
        if (*milliCpus != 0 && (*milliCpus < spec::kMinMilliCpus || *milliCpus > spec::kMaxMilliCpus)) {
            throw std::runtime_error("cpus must be between 0.1 and 1024 (0 = all cores)");
        }
        milliCpus_ = *milliCpus;
    }
    writeMeta();
}

void Container::create(const NetMode& net, const std::string& image, IdMode ids, const Options& o) {
    net_ = net;
    image_ = image;
    ids_ = ids;
    size_ = o.size.value_or(spec::kDefaultSizeBytes);
    relaxed_ = o.relaxed.value_or(false);
    memory_ = spec::kDefaultMemoryBytes;
    pids_ = spec::kDefaultPids;
    milliCpus_ = 0;
    if (o.memory || o.pids || o.milliCpus) setLimits(o.memory, o.pids, o.milliCpus);
    created_ = std::time(nullptr);
    fsutil::mkdirs(dir_);
    prepareLayout();
    writeMeta();
}

void Container::load() {
    readMeta();
    prepareLayout();
}

void Container::readMeta() {
    auto kv = parseKv(fsutil::readFile(metaFile()));
    net_ = Network::parse(kv.count("net") ? kv["net"] : "slirp");
    created_ = kv.count("created") ? std::stol(kv["created"]) : 0;
    image_ = kv.count("image") ? kv["image"] : "";
    ids_ = IdMap::parse(kv.count("ids") ? kv["ids"] : "single");
    size_ = kv.count("size") ? std::stoull(kv["size"]) : spec::kDefaultSizeBytes;
    memory_ = kv.count("memory") ? std::stoull(kv["memory"]) : spec::kDefaultMemoryBytes;
    pids_ = kv.count("pids") ? static_cast<std::uint32_t>(std::stoul(kv["pids"])) : spec::kDefaultPids;
    milliCpus_ = kv.count("cpus") ? static_cast<std::uint32_t>(std::stoul(kv["cpus"])) : 0;
    relaxed_ = kv.count("hardening") && kv["hardening"] == "relaxed";
}

std::optional<pid_t> Container::runningPid() const {
    try {
        std::istringstream in(fsutil::readFile(pidFile()));
        long pid = 0;
        std::string start;
        if (!(in >> pid >> start) || pid <= 1) return std::nullopt;
        auto now = startTime(static_cast<pid_t>(pid));
        if (now && *now == start) return static_cast<pid_t>(pid);
    } catch (const std::exception&) {
    }
    return std::nullopt;
}

void Container::recordPid(pid_t pid) const {
    auto st = startTime(pid);
    if (!st) throw std::runtime_error("container init exited before it could be recorded");
    fsutil::writeFile(pidFile(), std::to_string(pid) + " " + *st + "\n");
}

void Container::clearPid() const {
    std::error_code ec;
    fs::remove(pidFile(), ec);
}

Store::Store(fs::path root) : root_(std::move(root)) {
    std::string r = root_.string();
    if (r.find_first_of(",:") != std::string::npos) {
        throw std::runtime_error("state path must not contain ',' or ':' (overlayfs option syntax): " + r);
    }
    fsutil::mkdirs(containersDir(), 0700);
}

Store Store::fromEnvironment() {
    if (const char* h = std::getenv("DOOS_HOME"); h && *h) return Store(h);
    if (const char* x = std::getenv("XDG_DATA_HOME"); x && *x) return Store(fs::path(x) / "doos");
    const char* home = std::getenv("HOME");
    if (!home || !*home) throw std::runtime_error("HOME is not set; set DOOS_HOME");
    return Store(fs::path(home) / ".local/share/doos");
}

std::vector<Container> Store::list() const {
    std::vector<Container> out;
    for (const auto& e : fs::directory_iterator(containersDir())) {
        if (!e.is_directory() || e.is_symlink()) continue;
        std::string name = e.path().filename().string();
        try {
            Container::validateName(name);
            Container c(e.path(), name);
            c.readMeta();
            out.push_back(std::move(c));
        } catch (const std::exception& ex) {
            std::fprintf(stderr, "doos: skipping %s: %s\n", name.c_str(), ex.what());
        }
        if (out.size() >= spec::kMaxContainers) break;
    }
    std::sort(out.begin(), out.end(), [](const Container& a, const Container& b) { return a.name() < b.name(); });
    return out;
}

std::vector<std::string> Store::containersUsing(const std::string& image) const {
    std::vector<std::string> out;
    for (const auto& c : list()) {
        if (c.image() == image) out.push_back(c.name());
    }
    return out;
}

std::pair<Container, bool> Store::getOrCreate(const std::string& name, const Options& opts) {
    Container::validateName(name);
    fs::path dir = containersDir() / name;
    Container c(dir, name);
    std::error_code ec;
    std::optional<std::string> image;
    if (opts.image && !opts.image->empty()) image = ImageStore::normalize(*opts.image);

    auto verifyExisting = [&] {
        auto fixed = [&](const char* what, const std::string& have, const std::string& want) {
            if (have != want) {
                throw std::runtime_error("container '" + name + "' exists with " + what + "=" + have +
                                         "; this is fixed at creation");
            }
        };
        if (opts.net) fixed("net", Network::name(c.net()), Network::name(*opts.net));
        if (image) fixed("image", c.image().empty() ? "none" : c.image(), *image);
        if (opts.ids) fixed("ids", IdMap::name(c.ids()), IdMap::name(*opts.ids));
        if (opts.relaxed) fixed("hardening", c.relaxed() ? "relaxed" : "default", *opts.relaxed ? "relaxed" : "default");
        if (opts.size && *opts.size != c.sizeCap()) c.setSizeCap(*opts.size);
        if (opts.memory || opts.pids || opts.milliCpus) c.setLimits(opts.memory, opts.pids, opts.milliCpus);
    };

    if (fs::exists(dir / "meta", ec)) {
        FileLock lock(c.lockFile(), spec::kLockTimeoutMs);
        c.load();
        verifyExisting();
        return {c, false};
    }
    std::size_t count = 0;
    for (const auto& e : fs::directory_iterator(containersDir())) count += e.is_directory() ? 1 : 0;
    if (count >= spec::kMaxContainers) {
        throw std::runtime_error("container limit reached (" + std::to_string(spec::kMaxContainers) +
                                 "); delete one first");
    }

    // Resolve the image first (may download); ids follow the image so the layer ownership matches.
    IdMode ids = opts.ids.value_or(IdMap::best());
    if (image) {
        ImageStore store = images();
        if (!store.has(*image)) {
            if (!Catalog::knows(*image)) {
                throw std::runtime_error("unknown image '" + *image + "'. Add it: doos image add SOURCE [NAME]  (catalog: alpine, ubuntu, debian)");
            }
            std::fprintf(stderr, "doos: image '%s' not present, fetching it\n", image->c_str());
            store.add(Catalog::resolve(*image), ids);
        }
        IdMode imageIds = store.ids(*image);
        if (opts.ids && *opts.ids != imageIds) {
            throw std::runtime_error("image '" + *image + "' was extracted with ids=" + IdMap::name(imageIds) +
                                     "; use the same, or add the image again under another name");
        }
        ids = imageIds;
    }
    std::uint64_t size = opts.size.value_or(spec::kDefaultSizeBytes);
    if (size < spec::kMinSizeBytes || size > spec::kMaxSizeBytes) {
        throw std::runtime_error("size must be between " + formatSize(spec::kMinSizeBytes) + " and " +
                                 formatSize(spec::kMaxSizeBytes));
    }
    fsutil::mkdirs(dir, 0700);
    FileLock lock(c.lockFile(), spec::kLockTimeoutMs);
    if (fs::exists(dir / "meta", ec)) {  // lost a creation race
        c.load();
        verifyExisting();
        return {c, false};
    }
    Options resolved = opts;
    resolved.size = size;
    try {
        c.create(opts.net.value_or(NetMode::Slirp), image.value_or(""), ids, resolved);
    } catch (...) {
        // a container that failed halfway must not stay behind as an unreadable directory
        try {
            if (ids == IdMode::Subuid) UserNsJob::run(ids, [&] { fsutil::removeTree(dir); return 0; });
            else fsutil::removeTree(dir);
        } catch (...) {
        }
        throw;
    }
    if (c.net() == NetMode::Host) {
        std::fprintf(stderr, "doos: note: --net host shares the host's abstract unix sockets and loopback; "
                             "--net pasta gives the same reachability in a private namespace\n");
    }
    return {c, true};
}

std::string Store::remove(const std::string& name, bool confirmed) {
    Container::validateName(name);
    fs::path dir = containersDir() / name;
    std::error_code ec;
    auto st = fs::symlink_status(dir, ec);
    if (ec || !fs::is_directory(st)) throw std::runtime_error("no such container '" + name + "'");
    if (fs::canonical(dir).parent_path() != fs::canonical(containersDir())) {
        throw std::runtime_error("refusing to delete outside the container store: " + dir.string());
    }
    Container c(dir, name);
    FileLock lock(c.lockFile(), spec::kLockTimeoutMs);
    if (auto pid = c.runningPid()) {
        throw std::runtime_error("container '" + name + "' is running (pid " + std::to_string(*pid) + ")");
    }
    if (!confirmed) {
        std::size_t n = fsutil::countEntries(dir, 1000000);
        return "dry run: would delete " + dir.string() + " (" + std::to_string(n) + " entries). Add --yes to confirm.";
    }
    IdMode ids = IdMode::Single;
    try {
        c.readMeta();
        ids = c.ids();
    } catch (const std::exception&) {
        // damaged or half-created container: remove as plain user, files may be left behind loudly
    }
    if (ids == IdMode::Subuid) {
        // files owned by subuids can only be removed from inside the id map
        int rc = UserNsJob::run(ids, [&] { fsutil::removeTree(dir); return 0; });
        if (rc != 0) throw std::runtime_error("could not remove all files of '" + name + "'");
    } else {
        fsutil::removeTree(dir);
    }
    return "deleted " + name;
}

}  // namespace doos