// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include "RugnuxConsole.h"

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

#include <spdlog/details/os.h>
#include <spdlog/sinks/base_sink.h>

#ifdef _WIN32
#include <io.h>
#else
#include <unistd.h>
#endif

namespace {
    // Where the notes, warnings and the closing rule of a step start.
    const std::string INDENT(8, ' ');

    bool StdoutIsTerminal() {
#ifdef _WIN32
        return _isatty(_fileno(stdout)) != 0;
#else
        return isatty(fileno(stdout)) != 0;
#endif
    }

    // hh:mm:ss.mmm, the same clock and precision as the log file's timestamps.
    std::string TimeOfDay() {
        const auto now = std::chrono::system_clock::now();
        const std::tm tm = spdlog::details::os::localtime(std::chrono::system_clock::to_time_t(now));
        const auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(now.time_since_epoch()).count() % 1000;
        return fmt::format("{:02}:{:02}:{:02}.{:03}", tm.tm_hour, tm.tm_min, tm.tm_sec, ms);
    }

    std::string Duration(double seconds) {
        return fmt::format("{:.2f} s", seconds);
    }

    // Text set right-aligned at the console width after `left`, with at least two spaces between.
    std::string RightAligned(const std::string &left, const std::string &right) {
        const int pad = RugnuxConsole::WIDTH - static_cast<int>(left.size() + right.size());
        return left + std::string(std::max(2, pad), ' ') + right;
    }

    std::string Eta(double seconds) {
        if (seconds < 60.0)
            return fmt::format("{:.0f} s", seconds);
        const int s = static_cast<int>(seconds + 0.5);
        return fmt::format("{}:{:02} min", s / 60, s % 60);
    }

    class ConsoleSink : public spdlog::sinks::base_sink<std::mutex> {
        std::shared_ptr<RugnuxConsole> console;
    public:
        explicit ConsoleSink(std::shared_ptr<RugnuxConsole> c) : console(std::move(c)) {}
    protected:
        void sink_it_(const spdlog::details::log_msg &msg) override {
            if (msg.level >= spdlog::level::warn)
                console->Message(msg.level, std::string(msg.payload.data(), msg.payload.size()));
            else
                console->Heartbeat();
        }
        void flush_() override { std::fflush(stdout); }
    };
}

spdlog::sink_ptr MakeConsoleSink(std::shared_ptr<RugnuxConsole> console) {
    return std::make_shared<ConsoleSink>(std::move(console));
}

RugnuxConsole::RugnuxConsole(bool verbose) : verbose_(verbose), terminal_(StdoutIsTerminal()) {}

double RugnuxConsole::Elapsed() const {
    return Seconds(Clock::now() - start_);
}

void RugnuxConsole::LogLine(const std::string &text) {
    Logger("rugnux").Info("{}", text);
}

void RugnuxConsole::ClearStatus() {
    if (!status_shown_)
        return;
    std::fputs(("\r" + std::string(WIDTH, ' ') + "\r").c_str(), stdout);
    status_shown_ = false;
}

void RugnuxConsole::DrawStatus() {
    if (!terminal_ || verbose_ || status_.empty())
        return;
    std::string s = status_.substr(0, WIDTH);
    s.resize(WIDTH, ' ');
    std::fputs(("\r" + s).c_str(), stdout);
    std::fflush(stdout);
    status_shown_ = true;
}

void RugnuxConsole::Emit(const std::string &line) {
    ClearStatus();
    std::fputs(((terminal_ ? std::string() : TimeOfDay() + " ") + line + "\n").c_str(), stdout);
    DrawStatus();
    std::fflush(stdout);
}

void RugnuxConsole::Timed(const std::string &text, double seconds) {
    const std::string dur = Duration(seconds);
    if (text.size() + 2 + dur.size() <= static_cast<size_t>(WIDTH)) {
        Emit(RightAligned(text, dur));
        return;
    }
    // Too long for one line: wrapped under itself, the time on the last line.
    const size_t indent = text.find("... ") != std::string::npos ? text.find("... ") + 4 : 0;
    const size_t room = WIDTH - 2 - dur.size();
    std::string line = text.substr(0, indent);
    bool has_words = false;
    std::istringstream words(text.substr(indent));
    std::string word;
    while (words >> word) {
        if (has_words && line.size() + 1 + word.size() > room) {
            Emit(line);
            line = std::string(indent, ' ');
            has_words = false;
        }
        line += (has_words ? " " : "") + word;
        has_words = true;
    }
    Emit(RightAligned(line, dur));
}

void RugnuxConsole::Wrapped(const std::string &first, const std::string &rest, const std::string &text) {
    std::string line = first;
    bool line_has_words = false;
    std::istringstream paragraphs(text);
    std::string paragraph;
    while (std::getline(paragraphs, paragraph)) {
        std::istringstream words(paragraph);
        std::string word;
        while (words >> word) {
            if (line_has_words && line.size() + 1 + word.size() > static_cast<size_t>(WIDTH)) {
                Emit(line);
                line = rest;
                line_has_words = false;
            }
            line += (line_has_words ? " " : "") + word;
            line_has_words = true;
        }
        if (line_has_words) {
            Emit(line);
            line = rest;
            line_has_words = false;
        }
    }
}

void RugnuxConsole::Line(const std::string &text) {
    std::lock_guard lock(m_);
    Emit(text);
}

void RugnuxConsole::Rule(char c) {
    std::lock_guard lock(m_);
    Emit(std::string(WIDTH, c));
}

void RugnuxConsole::Field(const std::string &label, const std::string &value) {
    std::lock_guard lock(m_);
    const std::string first = fmt::format(" {:<14}", label);
    if (first.size() + value.size() <= static_cast<size_t>(WIDTH) && value.find('\n') == std::string::npos)
        Emit(first + value);   // as it is, so the spacing inside the value is kept
    else
        Wrapped(first, std::string(first.size(), ' '), value);
}

void RugnuxConsole::Close(const char *how) {
    if (step_ == 0)
        return;
    const auto now = Clock::now();
    const double total = Seconds(now - step_start_);
    const double rest = Seconds(now - mark_);
    std::string log_text = fmt::format("Step {} {} {}", step_, how, Duration(total));
    if (!verbose_) {
        status_.clear();
        // Time a step spent after its last note is shown where it is large, so the notes and the
        // total still add up for a reader looking for where the time went.
        if (notes_ > 0 && rest >= std::max(0.5, 0.1 * total))
            Timed(INDENT + "... (rest of the step)", rest);
        for (const auto &[kind, held] : warnings_held_) {
            std::string line = fmt::format("{}!! {} more like it, the last: {}", INDENT, held.first, held.second);
            if (line.size() > static_cast<size_t>(WIDTH))
                line = line.substr(0, WIDTH - 3) + "...";
            Emit(line);
        }
        const std::string label = std::string(" ") + how;
        const std::string dur = Duration(total);
        const int dashes = WIDTH - static_cast<int>(INDENT.size() + label.size() + 1 + dur.size());
        Emit(INDENT + std::string(std::max(4, dashes), '-') + label + " " + dur);
    }
    pending_log_.push_back(std::move(log_text));
    pending_timing_.push_back(fmt::format("TIMING step {} {:.3f}", step_title_, total));
    NvtxRangeEnd(step_nvtx_);
    warnings_held_.clear();
    step_ = 0;
}

void RugnuxConsole::OnStep(const std::string &title) {
    {
        std::lock_guard lock(m_);
        const int number = std::max(step_, last_step_) + 1;
        Close("done in");
        step_ = last_step_ = number;
        step_title_ = title;
        step_start_ = mark_ = Clock::now();
        step_nvtx_ = NvtxRangeStart(fmt::format("Step {} {}", step_, title));
        notes_ = 0;
        if (!verbose_)
            Emit(fmt::format("Step {:<2} {}", step_, title));
        pending_log_.push_back(fmt::format("Step {} {}", step_, title));
    }
    FlushLog();
}

void RugnuxConsole::OnNote(const std::string &text) {
    {
        std::lock_guard lock(m_);
        const auto now = Clock::now();
        const double dt = Seconds(now - mark_);
        if (!verbose_) {
            if (step_ == 0)
                Emit(INDENT + "... " + text);
            else
                Timed(INDENT + "... " + text, dt);
        }
        if (step_ > 0) {
            mark_ = now;
            ++notes_;
        }
        pending_log_.push_back(fmt::format("  ... {} ({})", text, Duration(dt)));
    }
    FlushLog();
}

void RugnuxConsole::OnPhase(const std::string &phase) {
    {
        std::lock_guard lock(m_);
        phase_ = phase;
        pending_log_.push_back("Phase: " + phase);
    }
    FlushLog();
    Heartbeat();
}

void RugnuxConsole::FlushLog() {
    std::vector<std::string> lines, timing;
    {
        std::lock_guard lock(m_);
        lines.swap(pending_log_);
        timing.swap(pending_timing_);
    }
    for (const auto &l : lines)
        LogLine(l);
    for (const auto &l : timing)
        Logger("rugnux").Debug("{}", l);
}

void RugnuxConsole::OnProgress(uint64_t done, uint64_t total) {
    if (verbose_ || total == 0)
        return;
    const auto now = Clock::now();
    const int64_t now_ns = std::chrono::duration_cast<std::chrono::nanoseconds>(now - start_).count();
    if (done < total && now_ns < next_draw_ns_.load())
        return;
    next_draw_ns_ = now_ns + (terminal_ ? 100'000'000 : 250'000'000);

    std::lock_guard lock(m_);
    // A new loop: another total, or the count starting over after the last loop ran to its end.
    // Workers report out of order, so "smaller than the last count" alone is not enough.
    if (total != loop_total_ || (loop_last_ == loop_total_ && done <= loop_total_ / 2)) {
        loop_total_ = total;
        loop_start_ = now;
        milestone_ = 0;
        loop_last_ = 0;
    }
    loop_last_ = std::max(loop_last_, done);
    const double elapsed = Seconds(now - loop_start_);
    const double rate = elapsed > 0.0 ? static_cast<double>(loop_last_) / elapsed : 0.0;

    if (terminal_) {
        if (loop_last_ >= loop_total_) {
            status_.clear();
            ClearStatus();
            return;
        }
        constexpr int BAR = 32;
        const int filled = static_cast<int>(BAR * loop_last_ / loop_total_);
        // The rate and what is left only once there is a second of loop to read them off.
        status_ = fmt::format("{}[{}{}] {:>6}/{}{}", INDENT, std::string(filled, '#'),
                              std::string(BAR - filled, '.'), loop_last_, loop_total_,
                              rate > 0.0 && elapsed > 1.0
                              ? fmt::format("  {:5.0f}/s  ETA {}", rate,
                                            Eta(static_cast<double>(loop_total_ - loop_last_) / rate))
                              : std::string());
        DrawStatus();
        return;
    }
    // Not a terminal: no bar, but a long pass says how far it has got once a quarter.
    const int quarter = static_cast<int>(4 * loop_last_ / loop_total_);
    while (milestone_ < std::min(quarter, 3)) {
        ++milestone_;
        if (elapsed > 5.0)
            Emit(fmt::format("{}... {}% ({} of {} images, {:.0f} per second)", INDENT, 25 * milestone_,
                             loop_last_, loop_total_, rate));
    }
}

void RugnuxConsole::Heartbeat() {
    if (!terminal_ || verbose_)
        return;
    const int64_t now_ns = std::chrono::duration_cast<std::chrono::nanoseconds>(Clock::now() - start_).count();
    if (now_ns < next_draw_ns_.load())
        return;
    // Never wait here: this runs inside every log call of every thread.
    std::unique_lock lock(m_, std::try_to_lock);
    if (!lock.owns_lock())
        return;
    next_draw_ns_ = now_ns + 250'000'000;
    const bool loop_running = loop_total_ > 0 && loop_last_ < loop_total_;
    if (step_ == 0 || loop_running)
        return;
    const double in_step = Seconds(Clock::now() - step_start_);
    if (in_step < 1.0)
        return;
    // A turning bar and what the pipeline says it is doing, so a long step is seen to be alive.
    static constexpr char SPIN[] = {'|', '/', '-', '\\'};
    const std::string time = fmt::format("{:.1f} s", in_step);
    std::string what = phase_.empty() ? std::string("working") : phase_;
    const size_t room = WIDTH - INDENT.size() - 4 - time.size() - 3;
    if (what.size() > room)
        what = what.substr(0, room - 3) + "...";
    status_ = RightAligned(INDENT + "[" + SPIN[spin_++ % 4] + "] " + what, time);
    DrawStatus();
}

void RugnuxConsole::Message(spdlog::level::level_enum level, const std::string &text) {
    std::lock_guard lock(m_);
    const std::string indent = step_ > 0 ? INDENT : std::string();
    const std::string tag = level >= spdlog::level::err ? "** ERROR: " : "!! ";
    // A warning the run has already shown in full, raised again by a later merge or pass with other
    // numbers in it, gets one line the first time it comes back and is then only counted, the count
    // shown when the step ends: it is still there, and the log has every one of them whole.
    std::string kind;
    for (char c : text.substr(0, 48))
        if (!std::isdigit(static_cast<unsigned char>(c)))
            kind += c;
    if (level < spdlog::level::err) {
        const int seen = warnings_seen_[kind]++;
        if (seen == 1) {
            std::string line = indent + tag + "again: " + text;
            if (line.size() > static_cast<size_t>(WIDTH))
                line = line.substr(0, WIDTH - 3) + "...";
            Emit(line);
        }
        if (seen >= 1) {
            if (seen >= 2) {
                ++warnings_held_[kind].first;
                warnings_held_[kind].second = text;
            }
            return;
        }
    }
    Wrapped(indent + tag, indent + std::string(tag.size(), ' '), text);
}

void RugnuxConsole::Finish() {
    {
        std::lock_guard lock(m_);
        Close("done in");
        status_.clear();
        ClearStatus();
    }
    FlushLog();
}

void RugnuxConsole::Fail() {
    {
        std::lock_guard lock(m_);
        Close("FAILED after");
        status_.clear();
        ClearStatus();
    }
    FlushLog();
}

bool RugnuxConsole::StepOpen() {
    std::lock_guard lock(m_);
    return step_ > 0;
}