#include "apps/ssh_app.h"

#include <Arduino.h>
#include <algorithm>
#include <SD.h>

#include "app_keys.h"
#include "text_wrap.h"
#include "ui/fonts.h"
#include "ui/theme.h"
#include "ui/widgets.h"

namespace roro {

namespace {
constexpr uint32_t kMessageMs = 4000;
constexpr uint32_t kFrameMs = 60;  // a busy terminal is drawn this often at most

std::string masked(size_t n) { return std::string(n, '*'); }

// The terminal keeps characters as the fonts have them, one byte each; drawing wants UTF-8.
void appendUtf8(std::string& out, uint8_t ch) {
    if (ch < 0x80) out += static_cast<char>(ch);
    else {
        out += static_cast<char>(0xC0 | (ch >> 6));
        out += static_cast<char>(0x80 | (ch & 0x3F));
    }
}

// Alt with ; and . looks back and forward, as in the Shell (Q258): lines to move, or 0.
int scrollStep(const KeyEvent& e) {
    if (!e.alt || e.ctrl) return 0;
    if (e.key == Key::Up || (e.key == Key::Char && e.ch == ';')) return 4;
    if (e.key == Key::Down || (e.key == Key::Char && e.ch == '.')) return -4;
    return 0;
}

uint16_t colour(int index) {
    uint8_t r, g, b;
    term::colourRgb(index, r, g, b);
    return lgfx::color565(r, g, b);
}
}  // namespace

// Ctrl with + and - step through these (Q257): from 60 columns by 20 rows to 26 by 8.
const SshApp::Font SshApp::kFonts[5] = {{&fonts::tiny, 4, 6}, {&fonts::small, 5, 8}, {&fonts::medium, 6, 10}, {&fonts::body, 6, 13}, {&fonts::large, 9, 15}};

void SshApp::grid(int& cols, int& rows) const {
    const auto& area = theme::kContent;
    cols = area.w / kFonts[font_].w;
    rows = area.h / kFonts[font_].h;
}

void SshApp::say(const std::string& text) {
    message_ = text;
    messageMs_ = millis();
    requestRedraw();
}

void SshApp::onEnter() {
    font_ = settings_.getInt(Setting::SshFont);
    hosts_ = term::SshHosts(settings_.getString(Setting::SshHosts)).list();
    list_.setCount(static_cast<int>(hosts_.size()) + 2);
    dialog_.reset();
    ask_ = Ask::None;
    message_.clear();
    scroll_ = 0;
    // A session that was left running is found again.
    SshService::State s = ssh_.state();
    view_ = s == SshService::State::Idle || (s == SshService::State::Ended && view_ != View::Session) ? View::Hosts : View::Session;
    shownState_ = SshService::State::Idle;
    requestRedraw();
}

void SshApp::connect(const term::SshTarget& target) {
    int cols, rows;
    grid(cols, rows);
    std::string why = ssh_.connect(target, cols, rows);
    if (!why.empty()) return say(why);
    password_ = LineEditor(96);
    scroll_ = 0;
    view_ = View::Session;
    shownState_ = SshService::State::Idle;
}

std::string SshApp::connectTo(const std::string& text) {
    term::SshTarget target;
    std::string why = term::parseSshTarget(text, target);
    if (!why.empty()) return why;
    if (ssh_.state() != SshService::State::Idle && ssh_.state() != SshService::State::Ended) return "a session is open already";
    font_ = settings_.getInt(Setting::SshFont);
    message_.clear();
    connect(target);
    return view_ == View::Session ? "" : message_;
}

void SshApp::setFont(int index) {
    index = std::clamp(index, 0, 4);
    if (index == font_) return;
    font_ = index;
    settings_.setInt(Setting::SshFont, font_);
    int cols, rows;
    grid(cols, rows);
    ssh_.resize(cols, rows);  // the far end is told, and redraws for the new size
    scroll_ = 0;
    say(std::to_string(cols) + " x " + std::to_string(rows));
}

void SshApp::writePublicKey() {
    std::string line = ssh_.publicKey() + "\n";
    storage_.runJob([line]() {
        if (!SD.exists("/ssh")) SD.mkdir("/ssh");
        File f = SD.open(kPublicKeyPath, FILE_WRITE);
        if (f) {
            f.write(reinterpret_cast<const uint8_t*>(line.data()), line.size());
            f.close();
        }
    });
}

void SshApp::help(std::vector<KeyHelp>& out) const {
    if (dialog_) return keys::add(out, keys::kDialog);
    switch (view_) {
        case View::Hosts: return keys::add(out, keys::kSsh);
        case View::Entry: return keys::add(out, keys::kText);
        case View::Key: return keys::add(out, keys::kSshKey);
        case View::Session:
            if (ssh_.state() == SshService::State::AskPassword) return keys::add(out, keys::kText);
            return keys::add(out, keys::kSshTerminal);
    }
}

const char* SshApp::helpTitle() const {
    switch (view_) {
        case View::Entry: return "SSH: a host";
        case View::Key: return "SSH: this device's key";
        case View::Session: return "SSH: the terminal";
        default: return nullptr;
    }
}

bool SshApp::sessionKey(const KeyEvent& e) {
    SshService::State state = ssh_.state();
    if (state == SshService::State::AskPassword) {
        switch (e.key) {
            case Key::Char: password_.insert(e.ch); break;
            case Key::Delete: password_.backspace(); break;
            case Key::Back:
                ssh_.disconnect();
                break;
            case Key::Select: {
                std::string typed = password_.text();
                password_ = LineEditor(96);
                ssh_.answerPassword(typed);
                break;
            }
            default: break;
        }
        return true;
    }
    if (state == SshService::State::Connecting) {
        if (e.key == Key::Back) ssh_.disconnect();
        return true;
    }
    if (state == SshService::State::Ended || state == SshService::State::Idle) {
        if (int step = scrollStep(e)) {  // what it said last can still be read
            int history = 0;
            ssh_.withTerminal([&](term::Terminal& t) { history = t.altScreen() ? 0 : t.historyCount(); });
            scroll_ = std::clamp(scroll_ + step, 0, history);
            return true;
        }
        ssh_.clear();  // read: its memory goes back
        view_ = View::Hosts;
        return true;
    }
    // Open: every key is the far end's, but these few.
    if (e.key == Key::Char && e.ctrl && e.alt && (e.ch == 'q' || e.ch == 'Q')) return ssh_.disconnect(), true;
    if (e.key == Key::Char && e.ctrl && (e.ch == '=' || e.ch == '+')) return setFont(font_ + 1), true;
    if (e.key == Key::Char && e.ctrl && (e.ch == '-' || e.ch == '_')) return setFont(font_ - 1), true;
    if (int step = scrollStep(e)) {
        int history = 0;
        ssh_.withTerminal([&](term::Terminal& t) { history = t.altScreen() ? 0 : t.historyCount(); });
        scroll_ = std::clamp(scroll_ + step, 0, history);
        return true;
    }
    bool app = false;
    ssh_.withTerminal([&](term::Terminal& t) { app = t.appCursorKeys(); });
    std::string bytes;
    switch (e.key) {
        case Key::Char: bytes = term::encodeKey(term::TermKey::Char, e.ch, e.ctrl, e.alt, app); break;
        case Key::Select: bytes = term::encodeKey(term::TermKey::Enter, 0, false, e.alt, app); break;
        case Key::Delete: bytes = term::encodeKey(term::TermKey::Backspace, 0, false, e.alt, app); break;
        case Key::Tab: bytes = term::encodeKey(term::TermKey::Tab, 0, false, false, app); break;
        // The key is printed "esc", and here that is what it is. With Alt it types the backtick it also carries.
        case Key::Back: bytes = e.alt ? "`" : term::encodeKey(term::TermKey::Escape, 0, false, false, app); break;
        case Key::Up: bytes = term::encodeKey(e.shift ? term::TermKey::PageUp : term::TermKey::Up, 0, false, false, app); break;
        case Key::Down: bytes = term::encodeKey(e.shift ? term::TermKey::PageDown : term::TermKey::Down, 0, false, false, app); break;
        case Key::Left: bytes = term::encodeKey(term::TermKey::Left, 0, false, false, app); break;
        case Key::Right: bytes = term::encodeKey(term::TermKey::Right, 0, false, false, app); break;
        default: break;
    }
    if (!bytes.empty()) {
        scroll_ = 0;
        ssh_.send(bytes);
    }
    return true;
}

bool SshApp::onKey(const KeyEvent& e) {
    requestRedraw();
    if (dialog_) {
        dialog_->onKey(e);
        int result = dialog_->result();
        if (result == DialogModel::kPending) return true;
        Ask asked = ask_;
        ask_ = Ask::None;
        dialog_.reset();
        if (asked == Ask::Trust) ssh_.answerTrust(result == 1);
        else if (asked == Ask::NewKey && result == 1) {
            std::string why = ssh_.makeKey();
            if (why.empty()) writePublicKey();
            say(why.empty() ? std::string("Made, and written to ") + kPublicKeyPath : why);
        } else if (asked == Ask::Forget && result == 1) {
            term::SshHosts saved(settings_.getString(Setting::SshHosts));
            term::SshTarget gone, other;
            term::parseSshTarget(hosts_[static_cast<size_t>(list_.selected())], gone);
            saved.remove(static_cast<size_t>(list_.selected()));
            settings_.setString(Setting::SshHosts, saved.stored());
            // Its fingerprint goes too, unless another saved host is the same server: met again, it is asked about again.
            bool shared = false;
            for (auto& h : saved.list()) shared = shared || (term::parseSshTarget(h, other).empty() && other.hostPort() == gone.hostPort());
            if (!shared) {
                term::SshKnownHosts known(settings_.getString(Setting::SshKnown));
                known.forget(gone.hostPort());
                settings_.setString(Setting::SshKnown, known.stored());
            }
            hosts_ = saved.list();
            list_.setCount(static_cast<int>(hosts_.size()) + 2);
        }
        return true;
    }
    switch (view_) {
        case View::Session: return sessionKey(e);
        case View::Entry:
            switch (e.key) {
                case Key::Char: entry_.insert(e.ch); break;
                case Key::Delete: entry_.backspace(); break;
                case Key::Left: entry_.left(); break;
                case Key::Right: entry_.right(); break;
                case Key::Back: view_ = View::Hosts; break;
                case Key::Select: {
                    term::SshTarget target;
                    std::string why = term::parseSshTarget(entry_.text(), target);
                    if (!why.empty()) say(why);
                    else connect(target);
                    break;
                }
                default: break;
            }
            return true;
        case View::Key:
            if (e.key == Key::Back) view_ = View::Hosts;
            else if (e.key == Key::Select || (e.key == Key::Char && (e.ch == 'n' || e.ch == 'N'))) {
                if (!ssh_.hasKey()) {
                    std::string why = ssh_.makeKey();
                    if (why.empty()) writePublicKey();
                    say(why.empty() ? std::string("Made, and written to ") + kPublicKeyPath : why);
                } else {
                    ask_ = Ask::NewKey;
                    dialog_.reset(new DialogModel({"Cancel", "New key"}));
                }
            } else if (e.key == Key::Char && (e.ch == 'w' || e.ch == 'W') && ssh_.hasKey()) {
                writePublicKey();
                say(std::string("Written to ") + kPublicKeyPath);
            }
            return true;
        case View::Hosts: {
            int count = static_cast<int>(hosts_.size());
            switch (e.key) {
                case Key::Up: list_.up(); break;
                case Key::Down: list_.down(); break;
                case Key::Back: return false;
                case Key::Select: {
                    int i = list_.selected();
                    if (i < count) {
                        term::SshTarget target;
                        if (term::parseSshTarget(hosts_[static_cast<size_t>(i)], target).empty()) connect(target);
                    } else if (i == count) {
                        entry_ = LineEditor(96);
                        view_ = View::Entry;
                    } else {
                        view_ = View::Key;
                    }
                    break;
                }
                case Key::Char:
                    if ((e.ch == 'd' || e.ch == 'D') && list_.selected() < count) {
                        ask_ = Ask::Forget;
                        dialog_.reset(new DialogModel({"Cancel", "Forget"}));
                    } else if (e.ch == 'n' || e.ch == 'N') {
                        entry_ = LineEditor(96);
                        view_ = View::Entry;
                    }
                    break;
                default: break;
            }
            return true;
        }
    }
    return true;
}

void SshApp::update(uint32_t nowMs) {
    if (!message_.empty() && nowMs - messageMs_ >= kMessageMs) {
        message_.clear();
        requestRedraw();
    }
    if (view_ != View::Session) return;
    SshService::State state = ssh_.state();
    if (state != shownState_) {
        shownState_ = state;
        if (state == SshService::State::AskTrust && !dialog_) {
            ask_ = Ask::Trust;
            dialog_.reset(new DialogModel({"Cancel", ssh_.remembered().empty() ? "Trust it" : "Replace"}));
        } else if (state != SshService::State::AskTrust && ask_ == Ask::Trust) {
            dialog_.reset();
            ask_ = Ask::None;
        }
        if (state == SshService::State::Open) {  // it got in: a host worth remembering (Q262)
            term::SshHosts saved(settings_.getString(Setting::SshHosts));
            saved.used(ssh_.target());
            settings_.setString(Setting::SshHosts, saved.stored());
            hosts_ = saved.list();
            list_.setCount(static_cast<int>(hosts_.size()) + 2);
            list_.select(0);
        }
        requestRedraw();
    }
    uint32_t revision = ssh_.revision();
    if (revision != shownRevision_ && nowMs - lastDrawMs_ >= kFrameMs) {
        shownRevision_ = revision;
        requestRedraw();
    }
}

void SshApp::drawHosts(Canvas& c) {
    const auto& area = theme::kContent;
    int count = static_cast<int>(hosts_.size());
    widgets::list(
        c, list_, {area.x, area.y, area.w, 8 * theme::kLineHeight},
        [&](int i) -> std::string { return i < count ? hosts_[static_cast<size_t>(i)] : i == count ? "New connection" : "This device's key"; },
        [&](int i) -> std::string { return i < count ? "" : i == count ? ">" : ssh_.hasKey() ? ">" : "none yet"; });
}

void SshApp::drawKey(Canvas& c) {
    const auto& area = theme::kContent;
    c.setFont(&fonts::body);
    c.setTextColor(theme::kMuted);
    if (!ssh_.hasKey()) {
        c.drawString("This device has no key yet.", 4, area.y + 4);
        c.drawString("With one, servers that know it", 4, area.y + 4 + 2 * theme::kLineHeight);
        c.drawString("ask for no password.", 4, area.y + 4 + 3 * theme::kLineHeight);
        c.setTextColor(theme::kText);
        c.drawString("Enter makes one.", 4, area.y + 4 + 5 * theme::kLineHeight);
        return;
    }
    c.drawString("Its public half, for a server's", 4, area.y + 2);
    c.drawString("authorized_keys file:", 4, area.y + 2 + theme::kLineHeight);
    c.setFont(&fonts::small);
    c.setTextColor(theme::kText);
    std::string pub = ssh_.publicKey();
    int y = area.y + 2 + 2 * theme::kLineHeight + 3;
    for (size_t at = 0; at < pub.size() && y < area.y + area.h - 30; at += 46, y += 9) c.drawString(pub.substr(at, 46).c_str(), 4, y);
    c.setFont(&fonts::small);
    c.setTextColor(theme::kMuted);
    c.drawString((std::string("It is also in ") + kPublicKeyPath + " (w writes it").c_str(), 4, area.y + area.h - 28);
    c.drawString("again). The private half never leaves.", 4, area.y + area.h - 19);
}

void SshApp::drawTerminal(Canvas& c) {
    const auto& area = theme::kContent;
    const Font& f = kFonts[font_];
    c.setFont(f.font);
    c.setTextDatum(top_left);
    ssh_.withTerminal([&](term::Terminal& t) {
        int rows = std::min(t.rows(), area.h / f.h), cols = std::min(t.cols(), area.w / f.w);
        int history = t.altScreen() ? 0 : t.historyCount(), back = std::min(scroll_, history);
        for (int r = 0; r < rows; r++) {
            int y = area.y + r * f.h, line = r - back;  // negative: a line of the history
            if (line < 0) {
                const std::string& text = t.historyLine(history + line);
                std::string shown;
                for (size_t i = 0; i < text.size() && static_cast<int>(i) < cols; i++) appendUtf8(shown, static_cast<uint8_t>(text[i]));
                c.setTextColor(theme::kMuted);
                c.drawString(shown.c_str(), area.x, y);
                continue;
            }
            for (int col = 0; col < cols;) {  // one run of the same colours at a time
                uint8_t attr = t.cell(line, col).attr;
                std::string run;
                int from = col;
                while (col < cols && t.cell(line, col).attr == attr) appendUtf8(run, t.cell(line, col++).ch);
                if (attr >> 4) c.fillRect(area.x + from * f.w, y, (col - from) * f.w, f.h, colour(attr >> 4));
                if (run.find_first_not_of(' ') == std::string::npos) continue;
                c.setTextColor(colour(attr & 15));
                c.drawString(run.c_str(), area.x + from * f.w, y);
            }
        }
        // The cursor: its cell with the colours swapped.
        int row = t.cursorRow() + back;
        if (t.cursorVisible() && row < rows && t.cursorCol() < cols && ssh_.state() == SshService::State::Open) {
            const term::Cell& cell = t.cell(t.cursorRow(), t.cursorCol());
            int x = area.x + t.cursorCol() * f.w, y = area.y + row * f.h;
            c.fillRect(x, y, f.w, f.h, colour(cell.attr & 15));
            std::string ch;
            appendUtf8(ch, cell.ch);
            c.setTextColor(colour(cell.attr >> 4));
            if (cell.ch != ' ') c.drawString(ch.c_str(), x, y);
        }
        if (back) {  // looking back: how far
            std::string where = "-" + std::to_string(back);
            c.setFont(&fonts::small);
            c.setTextDatum(top_right);
            c.fillRect(area.x + area.w - 5 * static_cast<int>(where.size()) - 3, area.y, 5 * static_cast<int>(where.size()) + 3, 9, theme::kAccent);
            c.setTextColor(0x0000);
            c.drawString(where.c_str(), area.x + area.w - 1, area.y + 1);
            c.setTextDatum(top_left);
        }
    });
}

void SshApp::drawSession(Canvas& c) {
    const auto& area = theme::kContent;
    SshService::State state = ssh_.state();
    if (state == SshService::State::Open || (state == SshService::State::Ended && ssh_.opened())) {
        drawTerminal(c);
        if (state == SshService::State::Ended) {
            std::string why = ssh_.status() + ". Any key.";
            if (why.size() > 47) why = ssh_.status();
            c.setFont(&fonts::small);
            c.fillRect(area.x, area.y + area.h - 11, area.w, 11, theme::kBackground);
            c.setTextColor(theme::kWarning);
            c.drawString(why.c_str(), 4, area.y + area.h - 9);
        }
        return;
    }
    c.setFont(&fonts::body);
    c.setTextColor(theme::kMuted);
    c.drawString(ssh_.target().text().c_str(), 4, area.y + 4);
    if (state == SshService::State::AskPassword) {
        std::string again = ssh_.status();
        c.setTextColor(again.empty() ? theme::kText : theme::kWarning);
        c.drawString(again.empty() ? "Password" : again.c_str(), 4, area.y + 4 + 2 * theme::kLineHeight);
        LineEditor stars(96);
        stars.setText(masked(password_.text().size()));
        widgets::lineEditor(c, stars, {4, area.y + 4 + 3 * theme::kLineHeight + 4, area.w - 8, 0});
        c.setFont(&fonts::small);
        c.setTextColor(theme::kMuted);
        c.drawString("It isn't kept. ` gives up.", 4, area.y + area.h - 10);
        return;
    }
    // Connecting, or it never got as far as a shell: what it is doing, or why not.
    bool ended = state == SshService::State::Ended;
    c.setTextColor(ended ? theme::kWarning : theme::kText);
    auto lines = wrapText(ssh_.status(), area.w - 8, widgets::bodyMeasure(c));
    for (size_t i = 0; i < lines.size() && i < 4; i++) c.drawString(lines[i].c_str(), 4, area.y + 4 + static_cast<int>(2 + i) * theme::kLineHeight);
    if (ended) {
        c.setFont(&fonts::small);
        c.setTextColor(theme::kMuted);
        c.drawString("Any key.", 4, area.y + area.h - 10);
    }
}

void SshApp::draw(Canvas& c) {
    lastDrawMs_ = millis();
    const auto& area = theme::kContent;
    c.setTextDatum(top_left);
    switch (view_) {
        case View::Hosts: drawHosts(c); break;
        case View::Entry:
            c.setFont(&fonts::body);
            c.setTextColor(theme::kMuted);
            c.drawString("Who, and where", 4, area.y + 4);
            widgets::lineEditor(c, entry_, {4, area.y + 22, area.w - 8, 0});
            c.drawString("user@host, or user@host:port", 4, area.y + 44);
            break;
        case View::Key: drawKey(c); break;
        case View::Session: drawSession(c); break;
    }
    if (!message_.empty() && !dialog_) {
        c.setFont(&fonts::small);
        c.fillRect(area.x, area.y + area.h - 11, area.w, 11, theme::kBackground);
        c.setTextColor(theme::kWarning);
        c.drawString(message_.c_str(), 4, area.y + area.h - 9);
    }
    if (dialog_ && ask_ == Ask::Trust) {
        std::string was = ssh_.remembered(), now = ssh_.fingerprint();
        // Three lines: one of words, and the fingerprint in two halves, since it is longer than a line.
        std::string shown = now.size() > 25 ? now.substr(0, 25) + " " + now.substr(25) : now;
        std::string host = ssh_.target().host.size() > 30 ? ssh_.target().host.substr(0, 28) + ".." : ssh_.target().host;
        if (was.empty()) widgets::dialog(c, "A server not met before", host + " " + shown, *dialog_);
        else widgets::dialog(c, "THE SERVER'S KEY CHANGED", "Someone in between? Now it is " + shown, *dialog_);
    } else if (dialog_ && ask_ == Ask::NewKey) {
        widgets::dialog(c, "A new key?", "Servers that know the old one will ask for a password again.", *dialog_);
    } else if (dialog_ && ask_ == Ask::Forget) {
        widgets::dialog(c, "Forget this host?", hosts_[static_cast<size_t>(std::min<int>(list_.selected(), static_cast<int>(hosts_.size()) - 1))], *dialog_);
    }
}

}  // namespace roro