#include "marlin2.h"
#include "esphome/core/log.h"
#include <string>

namespace esphome::marlin2 {
    static const char *TAG = "marlin2";

    #ifdef USE_SENSOR
    void Marlin2::add_sensor(const std::string &name, sensor::Sensor *sens) {
        sensors_.push_back({name, sens});
    }

    sensor::Sensor *Marlin2::find_sensor(const std::string &key) {
        for (const auto &pair : sensors_) {
            if (key == pair.first) {
                return pair.second;
            }
        }
        return nullptr;
    }
    #endif

    #ifdef USE_TEXT_SENSOR
    void Marlin2::add_text_sensor(const std::string &name, text_sensor::TextSensor *sens) {
        text_sensors_.push_back({name, sens});
    }

    text_sensor::TextSensor *Marlin2::find_text_sensor(const std::string &key) {
        for (const auto &pair : text_sensors_) {
            if (key == pair.first) {
                return pair.second;
            }
        }
        return nullptr;
    }
    #endif

    #ifdef USE_SELECT
    void Marlin2::add_select(const std::string &name, select::Select *sel) {
        selects_.push_back({name, sel});
    }

    select::Select *Marlin2::find_select(const std::string &key) {
        for (const auto &pair : selects_) {
            if (key == pair.first) {
                return pair.second;
            }
        }
        return nullptr;
    }
    #endif

    #ifdef USE_BINARY_SENSOR
    void Marlin2::add_binary_sensor(const std::string &name, binary_sensor::BinarySensor *bs) {
        binary_sensors_.push_back({name, bs});
    }

    binary_sensor::BinarySensor *Marlin2::find_binary_sensor(const std::string &key) {
        for (const auto &pair : binary_sensors_) {
            if (key == pair.first) {
                return pair.second;
            }
        }
        return nullptr;
    }
    #endif

    void Marlin2::setup() {
        marlin_output_.reserve(256);
        marlin_response_output_.reserve(256);
        marlin_time_.reserve(32);
        printer_state_.reserve(32);

        write_str("\r\n\r\nM155 S10\r\n");
        write_str("\r\n\r\nM117 ESP Home Connected!\r\n");
        flush();

        set_printer_state("IDLE");
    }

    void Marlin2::dump_config() {
        ESP_LOGCONFIG(TAG, "Marlin2:");
        this->check_uart_settings(115200);
        LOG_UPDATE_INTERVAL(this);
    }

    void Marlin2::write(std::string gcode) {
        ESP_LOGD(TAG, "->GCODE: %s", gcode.c_str());
        write_str((std::string("\r\n\r\n") + gcode + std::string("\r\n")).c_str());
        flush();
    }

    void Marlin2::update() {
        while (available()) {
            char c = read();
            if (c == '\n' || c == '\r') {
                process_line();
            } else {
                marlin_output_ += c;
            }
        }

        if (millis() - millis_progress_ > 15000 && print_progress_ != 100) {
            millis_progress_ = millis();
            write_str("M27\r\nM31\r\n");
        }
    }

    void Marlin2::process_line() {
        if (marlin_output_.size() < 3) {
            marlin_output_ = "";
            return;
        }

        if (marlin_output_.compare("ok") == 0 || marlin_output_.compare(" ok") == 0) {
            listing_file_ = false;
            marlin_output_ = "";
            return;
        }

        ESP_LOGD(TAG, "DEBUG>#%s#", marlin_output_.c_str());

        if (!listing_file_ && marlin_output_.compare("Begin file list") == 0) {
            ESP_LOGD(TAG, "Listing of files started!");
            listing_file_ = true;
            file_table_.clear();
            marlin_output_ = "";
            return;
        }

        if (listing_file_ && marlin_output_.compare("End file list") == 0) {
            ESP_LOGD(TAG, "Listing of files stopped!");
            listing_file_ = false;

            #ifdef USE_BINARY_SENSOR
            if (find_binary_sensor("sd_card_present") != nullptr) {
                find_binary_sensor("sd_card_present")->publish_state(true);
            }
            #endif

            publish_sd_files();

            #ifdef USE_SENSOR
            if (find_sensor("sd_card_file_count") != nullptr) {
                find_sensor("sd_card_file_count")->publish_state(static_cast<float>(file_table_.size()));
            }
            #endif

            marlin_output_ = "";
            return;
        }

        if (listing_file_ && marlin_output_.find(".GCO ") > 1) {
            size_t first_space = marlin_output_.find(' ');
            size_t second_space = marlin_output_.find(' ', first_space + 1);

            if (first_space != std::string::npos) {
                std::string short_name = marlin_output_.substr(0, first_space);

                if (short_name.find("/") != std::string::npos) {
                    ESP_LOGD(TAG, "Skipping subdirectory: %s", short_name.c_str());
                    marlin_output_ = "";
                    return;
                }

                std::string long_name = short_name;

                if (second_space != std::string::npos && second_space > first_space) {
                    long_name = marlin_output_.substr(second_space + 1);
                }

                file_table_[long_name] = short_name;
            }

            marlin_output_ = "";
            return;
        }

        // Parse periodic temperature read out message
        if (
            marlin_output_.find(" T:")    == 0 ||
            marlin_output_.find("T:")     == 0 ||
            marlin_output_.find("ok T:")  == 0 ||
            marlin_output_.find(" ok T:") == 0
        ) {
            float ext_temperature, ext_set_temperature, bed_temperature, bed_set_temperature;
            if (process_temp_msg(&ext_temperature, &ext_set_temperature, &bed_temperature, &bed_set_temperature) != 0) {
                #ifdef USE_SENSOR
                if (find_sensor("bed_temperature") != nullptr)
                    find_sensor("bed_temperature")->publish_state(bed_temperature);

                if (find_sensor("bed_set_temperature") != nullptr)
                    find_sensor("bed_set_temperature")->publish_state(bed_set_temperature);

                if (find_sensor("ext_temperature") != nullptr)
                    find_sensor("ext_temperature")->publish_state(ext_temperature);

                if (find_sensor("ext_set_temperature") != nullptr)
                    find_sensor("ext_set_temperature")->publish_state(ext_set_temperature);
                #endif
                #ifdef USE_TEXT_SENSOR
                if (bed_set_temperature == 0.0 && ext_set_temperature == 0.0) {
                    if (ext_temperature < 32.0 && bed_temperature < 32.0) {
                        set_printer_state("IDLE");
                    } else if (ext_temperature < 150.0 && bed_temperature < 55.0) {
                        set_printer_state("COOLING");
                    }
                }
                if (print_progress_ == 0.0 && (bed_set_temperature != 0.0 || ext_set_temperature != 0.0)) {
                    set_printer_state("PREHEATING");
                }
                #endif
            }

            marlin_output_ = "";
            return;
        }

        // Parse progress of the print
        if (marlin_output_.find("SD printing byte") == 0) {
            print_progress_ = process_progress_msg();
            #ifdef USE_SENSOR
            if (find_sensor("print_progress") != nullptr)
                find_sensor("print_progress")->publish_state(print_progress_);
            #endif

            set_printer_state("PRINTING");
            marlin_output_ = "";
            return;
        }

        // Parse print time
        if (marlin_output_.find("echo:Print time: ") == 0) {
            double current = 0;
            double remaining = 0;
            if (process_print_time_msg(&current, &remaining, print_progress_) != 0) {
                #ifdef USE_SENSOR
                if (find_sensor("print_time") != nullptr)
                    find_sensor("print_time")->publish_state(current);

                if (find_sensor("print_time_remaining") != nullptr)
                    find_sensor("print_time_remaining")->publish_state(remaining);
                #endif
            }

            marlin_output_ = "";
            return;
        }

        // File opened: SALMA_~2.GCO Size: 12279971
        if (marlin_output_.find("File opened: ") == 0) {
            size_t start_pos = std::strlen("File opened: ");
            size_t size_pos  = marlin_output_.find(" Size: ");

            if (size_pos != std::string::npos) {
                std::string filename = from_dos_name(marlin_output_.substr(start_pos, size_pos - start_pos));
                ESP_LOGD(TAG, "Soubor: %s", filename.c_str());

                #ifdef USE_TEXT_SENSOR
                if (find_text_sensor("sd_card_file_selected") != nullptr) {
                    find_text_sensor("sd_card_file_selected")->publish_state(filename.c_str());
                }
                #endif
            }

            marlin_output_ = "";
            return;
        }

        // Print from SD card started
        if (marlin_output_.compare("File selected") == 0) {
            set_printer_state("PRINTING");
            marlin_output_ = "";
            return;
        }

        // Print finished
        if (marlin_output_.compare("Done printing") == 0) {
            print_progress_ = 100;
            #ifdef USE_SENSOR
            if (find_sensor("print_progress") != nullptr)
                find_sensor("print_progress")->publish_state(print_progress_);

            if (find_sensor("print_time_remaining") != nullptr)
                find_sensor("print_time_remaining")->publish_state(0);
            #endif
            #ifdef USE_TEXT_SENSOR
            set_printer_state("FINISHED");
            #endif

            marlin_output_ = "";
            return;
        }

        // Print aborted
        if (marlin_output_.compare("Print Aborted") == 0) {
            set_printer_state("STOPPED");
            marlin_output_ = "";
            return;
        }

        ESP_LOGD(TAG, ">#%s#", marlin_output_.c_str());
        marlin_output_ = "";
    }

    void Marlin2::publish_sd_files() {
        #if defined(USE_TEXT_SENSOR) || defined(USE_SELECT)
        std::string chunk;
        uint8_t count = 0;

        #ifdef USE_SELECT
        std::vector<std::string> file_names;
        #endif

        for (const auto &[key, value] : file_table_) {
            if (count >= max_sd_files_) {
                ESP_LOGW(TAG, "SD file list truncated at %u files (max_sd_files=%u, total=%u)",
                         max_sd_files_, max_sd_files_, (unsigned) file_table_.size());
                break;
            }

            #ifdef USE_SELECT
            file_names.push_back(key);
            #endif

            if (!chunk.empty()) chunk += "|";
            chunk += key;
            count++;
        }

        #ifdef USE_TEXT_SENSOR
        if (find_text_sensor("sd_card_files") != nullptr) {
            find_text_sensor("sd_card_files")->publish_state(chunk);
        }
        #endif

        #ifdef USE_SELECT
        if (find_select("sd_card_file_select") != nullptr) {
            find_select("sd_card_file_select")->traits.set_options(file_names);
        }
        #endif
        #endif
    }

    int Marlin2::process_temp_msg(float *ext_temperature, float *ext_set_temperature, float *bed_temperature, float *bed_set_temperature) {
        float dc;

        while (marlin_output_.find(" ") != std::string::npos)
            marlin_output_.erase(marlin_output_.find(' '), 1);

        while (marlin_output_.find("ok") != std::string::npos)
            marlin_output_.erase(marlin_output_.find("ok"), 2);

        if (sscanf(marlin_output_.c_str(), "T:%f/%fB:%f/%f", ext_temperature, ext_set_temperature, bed_temperature, bed_set_temperature) == 4)
            return 1;

        if (sscanf(marlin_output_.c_str(), "T:%f/%f(%f)B:%f/%f(%f)", ext_temperature, ext_set_temperature, &dc, bed_temperature, bed_set_temperature, &dc) == 6)
            return 2;

        if (sscanf(marlin_output_.c_str(), "T:%f/%fT0:%f/%fT1:%f/%fB:%f/%f", ext_temperature, ext_set_temperature, &dc, &dc, &dc, &dc, bed_temperature, bed_set_temperature) == 8)
            return 3;

        if (sscanf(marlin_output_.c_str(), "T0:%f/%fT1:%f/%fB:%f/%f", ext_temperature, ext_set_temperature, &dc, &dc, bed_temperature, bed_set_temperature) == 6)
            return 4;

        return 0;
    }

    float Marlin2::process_progress_msg() {
        size_t slash_pos = marlin_output_.find('/');
        if (slash_pos == std::string::npos || slash_pos + 1 >= marlin_output_.size())
            return 0.0f;

        float current = std::stof(marlin_output_.substr(17));
        float total   = std::stof(marlin_output_.substr(slash_pos + 1));

        if (total < 1.0f)
            return 0.0f;

        return roundf((current * 100.0f) / total);
    }

    int Marlin2::process_print_time_msg(double *current, double *remaining, float progress) {
        marlin_time_ = marlin_output_.substr(16);
        float d = 0, h = 0, m = 0, s = 0;

        if (sscanf(marlin_time_.c_str(), "%fd %fh %fm %fs", &d, &h, &m, &s) != 4) {
            d = 0;
            if (sscanf(marlin_time_.c_str(), "%fh %fm %fs", &h, &m, &s) != 3) {
                d = 0; h = 0;
                if (sscanf(marlin_time_.c_str(), "%fm %fs", &m, &s) != 2) {
                    d = 0; h = 0; m = 0;
                    if (sscanf(marlin_time_.c_str(), "%fs", &s) != 1) {
                        return 0;
                    }
                }
            }
        }

        *current = round(((d) * 24 * 60 * 60) + ((h) * 60 * 60) + ((m) * 60) + (s));

        if (progress != 0.0 && progress != 100.0) {
            *remaining = (((100 * *current) / round(progress)) - *current);
        }

        return 1;
    }

    void Marlin2::set_printer_state(std::string status) {
        #ifdef USE_TEXT_SENSOR
        if (find_text_sensor("printer_state") != nullptr) {
            find_text_sensor("printer_state")->publish_state(status);
        }
        #endif
    }

    std::string Marlin2::to_dos_name(std::string filename) {
        auto it = file_table_.find(filename);
        if (it != file_table_.end()) {
            return it->second;
        }
        return filename;
    }

    std::string Marlin2::from_dos_name(std::string dos_filename) {
        for (const auto &[key, value] : file_table_) {
            if (value == dos_filename) {
                return key;
            }
        }
        return dos_filename;
    }

}  // namespace esphome::marlin2