/*
 * uuid-telnet - Telnet service
 * Copyright 2019  Simon Arlott
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include "uuid/telnet.h"

#include <Arduino.h>

#include <algorithm>
#include <string>
#include <vector>

namespace uuid {

namespace telnet {

TelnetStream::TelnetStream(WiFiClient & client)
    : client_(client) {
    output_buffer_.reserve(BUFFER_SIZE);
}

void TelnetStream::start() {
    raw_write({IAC, WILL, OPT_ECHO, IAC, WILL, OPT_BINARY, IAC, WILL, OPT_SGA, IAC, DONT, OPT_ECHO, IAC, DO, OPT_BINARY, IAC, DO, OPT_SGA});
}

int TelnetStream::available() {
    if (peek() == -1) {
        return 0;
    } else {
        return 1;
    }
}

int TelnetStream::read() {
    if (peek_ != -1) {
        int data = peek_;
        peek_    = -1;
        return data;
    }

    buffer_flush();

restart:
    int data = raw_read();

    if (data == -1) {
        return -1;
    }

    unsigned char c = data;

    if (sub_negotiation_) {
        if (previous_raw_in_ == IAC) {
            switch (c) {
            case SE:
                sub_negotiation_ = false;
                previous_raw_in_ = 0;
                goto restart;

            case IAC:
                previous_raw_in_ = 0;
                goto restart;
            }
        } else {
            switch (c) {
            case IAC:
                previous_raw_in_ = c;
                goto restart;

            default:
                previous_raw_in_ = 0;
                goto restart;
            }
        }
    } else {
        if (previous_raw_in_ == IAC) {
            switch (c) {
            case IP:
                // Interrupt (^C)
                previous_raw_in_ = 0;
                c                = '\x03';
                break;

            case EC:
                // Backspace (^H)
                previous_raw_in_ = 0;
                c                = '\x08';
                break;

            case EL:
                // Delete line (^U)
                previous_raw_in_ = 0;
                c                = '\x15';
                break;

            case IAC:
                previous_raw_in_ = 0;
                break;

            case SB:
            case WILL:
            case WONT:
            case DO:
            case DONT:
                previous_raw_in_ = c;
                goto restart;

            case SE:
            case DM:
            case BRK:
            case AO:
            case AYT:
            case GA:
            case NOP:
            default:
                previous_raw_in_ = 0;
                goto restart;
            }
        } else if (previous_raw_in_ == SB) {
            sub_negotiation_ = true;
            previous_raw_in_ = 0;
            goto restart;
        } else if (previous_raw_in_ == WILL || previous_raw_in_ == WONT) {
            switch (c) {
            case OPT_ECHO:
                // Don't do these
                raw_write({IAC, DONT, c});
                break;

            case OPT_BINARY:
            case OPT_SGA:
                // Do these
                raw_write({IAC, DO, c});
                break;

            default:
                // Don't do anything else
                raw_write({IAC, DONT, c});
                break;
            }

            previous_raw_in_ = 0;
            goto restart;
        } else if (previous_raw_in_ == DO) {
            switch (c) {
            case OPT_ECHO:
            case OPT_BINARY:
            case OPT_SGA:
                // These are always enabled
                break;

            default:
                // Refuse to do anything else
                raw_write({IAC, WONT, c});
                break;
            }

            previous_raw_in_ = 0;
            goto restart;
        } else if (previous_raw_in_ == DONT) {
            switch (c) {
            case OPT_ECHO:
            case OPT_BINARY:
            case OPT_SGA:
                // Insist that we do these
                raw_write({IAC, WILL, c});
                break;

            default:
                // Everything else is always disabled
                break;
            }

            previous_raw_in_ = 0;
            goto restart;
        } else {
            switch (c) {
            case IAC:
                previous_raw_in_ = c;
                goto restart;

            default:
                previous_raw_in_ = 0;
                break;
            }
        }
    }

    if (previous_in_ == CR) {
        if (c == NUL) {
            previous_in_ = 0;
            goto restart;
        }
    }

    previous_in_ = c;
    return c;
}

int TelnetStream::peek() {
    buffer_flush();

    // It's too complicated to implement this by calling peek()
    // on the original stream, especially if the original stream
    // doesn't actually support peeking.
    if (peek_ == -1) {
        peek_ = read();
    }

    return peek_;
}

size_t TelnetStream::write(uint8_t data) {
    if (previous_out_ == CR && data != LF) {
        previous_out_ = data;

        if (raw_write({NUL, data}) != 2) {
            return 0;
        }
    } else {
        previous_out_ = data;
    }

    if (data == IAC) {
        if (raw_write({IAC, IAC}) != 2) {
            return 0;
        }
    } else {
        if (raw_write(data) != 1) {
            return 0;
        }
    }

    return 1;
}

size_t TelnetStream::write(const uint8_t * buffer, size_t size) {
    std::vector<unsigned char> data;
    data.reserve(size);

    while (size-- > 0) {
        unsigned char c = *buffer++;

        if (previous_out_ == CR && c != LF) {
            data.push_back((unsigned char)NUL);
        }

        if (c == IAC) {
            data.push_back((unsigned char)IAC);
        }

        previous_out_ = c;
        data.push_back(c);
    }

    size_t len = raw_write(data);
    if (len < size) {
        len = 0;
    }
    return len;
}

void TelnetStream::flush() {
    // This is a pure virtual function in Arduino's Stream class, which
    // makes no sense because that class is for input and this is an
    // output function. Later versions move it to Print as an empty
    // virtual function so this is here for backward compatibility.
}

int TelnetStream::raw_available() {
    return client_.available();
}

int TelnetStream::raw_read() {
    return client_.read();
}

void TelnetStream::buffer_flush() {
    if (!output_buffer_.empty()) {
        size_t len = client_.write(reinterpret_cast<const unsigned char *>(output_buffer_.data()), output_buffer_.size());
        if (len != output_buffer_.size()) {
            client_.stop();
        }
        output_buffer_.clear();
        output_buffer_.shrink_to_fit();
    }
}

size_t TelnetStream::raw_write(unsigned char data) {
    output_buffer_.push_back(data);

    if (output_buffer_.size() >= BUFFER_SIZE) {
        buffer_flush();
    }

    return 1;
}

size_t TelnetStream::raw_write(const std::vector<unsigned char> & data) {
    return raw_write(reinterpret_cast<const unsigned char *>(data.data()), data.size());
}

size_t TelnetStream::raw_write(const uint8_t * buffer, size_t size) {
    size_t offset    = 0;
    size_t remaining = size;

    if (!output_buffer_.empty()) {
        // Fill the rest of the buffer
        size_t block = std::min(remaining, BUFFER_SIZE - output_buffer_.size());

        output_buffer_.insert(output_buffer_.end(), buffer, buffer + block);
        offset += block;
        remaining -= block;

        if (output_buffer_.size() >= BUFFER_SIZE) {
            buffer_flush();
        }
    }

    if (remaining >= BUFFER_SIZE) {
        // Output directly if it won't fit in the buffer
        size_t len = client_.write(buffer + offset, remaining);
        if (len != remaining) {
            client_.stop();
            return offset + len;
        }
    } else if (remaining > 0) {
        // Put the rest in the buffer
        output_buffer_.insert(output_buffer_.end(), buffer + offset, buffer + offset + remaining);
    }

    return size;
}

} // namespace telnet

} // namespace uuid