#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <algorithm>
#include <cmath>
#include <fstream>
#include <stdexcept>
#include <string>
#include <vector>

#include <opencv2/opencv.hpp>

#include "rknn_api.h"

struct TextBox {
    cv::Rect rect;
    float score;
};

static std::vector<unsigned char> read_file_binary(const std::string& path) {
    std::ifstream file(path, std::ios::binary | std::ios::ate);
    if (!file) {
        throw std::runtime_error("failed to open file: " + path);
    }

    std::streamsize size = file.tellg();
    file.seekg(0, std::ios::beg);

    std::vector<unsigned char> data(size);
    if (!file.read(reinterpret_cast<char*>(data.data()), size)) {
        throw std::runtime_error("failed to read file: " + path);
    }

    return data;
}

static std::vector<std::string> load_dict(const std::string& path) {
    std::ifstream ifs(path);
    if (!ifs) {
        throw std::runtime_error("failed to open dict: " + path);
    }

    std::vector<std::string> dict;
    std::string line;

    while (std::getline(ifs, line)) {
        if (!line.empty() && line.back() == '\r') {
            line.pop_back();
        }
        dict.push_back(line);
    }

    return dict;
}

class RknnModel {
public:
    explicit RknnModel(const std::string& model_path) {
        model_data_ = read_file_binary(model_path);

        int ret = rknn_init(&ctx_, model_data_.data(), model_data_.size(), 0, NULL);
        if (ret != RKNN_SUCC) {
            throw std::runtime_error("rknn_init failed: " + model_path + ", ret=" + std::to_string(ret));
        }

        memset(&io_num_, 0, sizeof(io_num_));
        ret = rknn_query(ctx_, RKNN_QUERY_IN_OUT_NUM, &io_num_, sizeof(io_num_));
        if (ret != RKNN_SUCC) {
            throw std::runtime_error("rknn_query IN_OUT_NUM failed, ret=" + std::to_string(ret));
        }

        input_attrs_.resize(io_num_.n_input);
        output_attrs_.resize(io_num_.n_output);

        for (uint32_t i = 0; i < io_num_.n_input; ++i) {
            memset(&input_attrs_[i], 0, sizeof(rknn_tensor_attr));
            input_attrs_[i].index = i;

            ret = rknn_query(ctx_, RKNN_QUERY_INPUT_ATTR, &input_attrs_[i], sizeof(rknn_tensor_attr));
            if (ret != RKNN_SUCC) {
                throw std::runtime_error("rknn_query INPUT_ATTR failed, ret=" + std::to_string(ret));
            }
        }

        for (uint32_t i = 0; i < io_num_.n_output; ++i) {
            memset(&output_attrs_[i], 0, sizeof(rknn_tensor_attr));
            output_attrs_[i].index = i;

            ret = rknn_query(ctx_, RKNN_QUERY_OUTPUT_ATTR, &output_attrs_[i], sizeof(rknn_tensor_attr));
            if (ret != RKNN_SUCC) {
                throw std::runtime_error("rknn_query OUTPUT_ATTR failed, ret=" + std::to_string(ret));
            }
        }
    }

    ~RknnModel() {
        if (ctx_ != 0) {
            rknn_destroy(ctx_);
            ctx_ = 0;
        }
    }

    const rknn_tensor_attr& input_attr(int index = 0) const {
        return input_attrs_.at(index);
    }

    const rknn_tensor_attr& output_attr(int index = 0) const {
        return output_attrs_.at(index);
    }

    std::vector<float> run_uint8_nhwc(const unsigned char* input_data,
                                      int input_size,
                                      int output_index = 0) {
        rknn_input input;
        memset(&input, 0, sizeof(input));

        input.index = 0;
        input.buf = const_cast<unsigned char*>(input_data);
        input.size = input_size;
        input.pass_through = 0;
        input.type = RKNN_TENSOR_UINT8;
        input.fmt = RKNN_TENSOR_NHWC;

        int ret = rknn_inputs_set(ctx_, 1, &input);
        if (ret != RKNN_SUCC) {
            throw std::runtime_error("rknn_inputs_set failed, ret=" + std::to_string(ret));
        }

        ret = rknn_run(ctx_, NULL);
        if (ret != RKNN_SUCC) {
            throw std::runtime_error("rknn_run failed, ret=" + std::to_string(ret));
        }

        std::vector<rknn_output> outputs(io_num_.n_output);
        for (uint32_t i = 0; i < io_num_.n_output; ++i) {
            memset(&outputs[i], 0, sizeof(rknn_output));
            outputs[i].index = i;
            outputs[i].want_float = 1;
        }

        ret = rknn_outputs_get(ctx_, io_num_.n_output, outputs.data(), NULL);
        if (ret != RKNN_SUCC) {
            throw std::runtime_error("rknn_outputs_get failed, ret=" + std::to_string(ret));
        }

        const rknn_tensor_attr& out_attr = output_attrs_.at(output_index);
        int n = out_attr.n_elems;

        std::vector<float> result(n);
        memcpy(result.data(), outputs[output_index].buf, n * sizeof(float));

        rknn_outputs_release(ctx_, io_num_.n_output, outputs.data());

        return result;
    }

private:
    rknn_context ctx_ = 0;
    rknn_input_output_num io_num_;
    std::vector<rknn_tensor_attr> input_attrs_;
    std::vector<rknn_tensor_attr> output_attrs_;
    std::vector<unsigned char> model_data_;
};

static void print_tensor_shape(const char* tag, const rknn_tensor_attr& attr) {
    printf("%s: ", tag);
    for (int i = 0; i < attr.n_dims; ++i) {
        printf("%d", attr.dims[i]);
        if (i + 1 < attr.n_dims) {
            printf(" x ");
        }
    }
    printf(", fmt=%d, type=%d, size=%d\n", attr.fmt, attr.type, attr.size);
}

static cv::Mat resize_with_padding(const cv::Mat& src,
                                   int dst_w,
                                   int dst_h,
                                   float& scale,
                                   int& pad_x,
                                   int& pad_y) {
    int src_w = src.cols;
    int src_h = src.rows;

    scale = std::min(static_cast<float>(dst_w) / src_w,
                     static_cast<float>(dst_h) / src_h);

    int new_w = std::max(1, static_cast<int>(std::round(src_w * scale)));
    int new_h = std::max(1, static_cast<int>(std::round(src_h * scale)));

    pad_x = 0;
    pad_y = 0;

    cv::Mat resized;
    cv::resize(src, resized, cv::Size(new_w, new_h), 0, 0, cv::INTER_LINEAR);

    cv::Mat padded(dst_h, dst_w, src.type(), cv::Scalar(0, 0, 0));
    resized.copyTo(padded(cv::Rect(pad_x, pad_y, new_w, new_h)));

    return padded;
}

static std::vector<TextBox> postprocess_det(const std::vector<float>& prob,
                                            int map_w,
                                            int map_h,
                                            int orig_w,
                                            int orig_h,
                                            float scale,
                                            int pad_x,
                                            int pad_y) {
    cv::Mat prob_map(map_h, map_w, CV_32FC1, const_cast<float*>(prob.data()));
    cv::Mat binary;

    cv::threshold(prob_map, binary, 0.30, 255.0, cv::THRESH_BINARY);
    binary.convertTo(binary, CV_8UC1);

    cv::morphologyEx(binary, binary, cv::MORPH_CLOSE, cv::Mat(), cv::Point(-1, -1), 1);

    std::vector<std::vector<cv::Point> > contours;
    cv::findContours(binary, contours, cv::RETR_LIST, cv::CHAIN_APPROX_SIMPLE);

    std::vector<TextBox> boxes;

    for (size_t i = 0; i < contours.size(); ++i) {
        double area = cv::contourArea(contours[i]);
        if (area < 20.0) {
            continue;
        }

        cv::Rect r = cv::boundingRect(contours[i]);

        float score_sum = 0.0f;
        int score_count = 0;

        for (int y = r.y; y < r.y + r.height; ++y) {
            const float* row = prob_map.ptr<float>(y);
            for (int x = r.x; x < r.x + r.width; ++x) {
                if (binary.at<unsigned char>(y, x) > 0) {
                    score_sum += row[x];
                    score_count++;
                }
            }
        }

        float score = score_count > 0 ? score_sum / score_count : 0.0f;
        if (score < 0.35f) {
            continue;
        }

        int margin = std::max(2, static_cast<int>(std::round(std::max(r.width, r.height) * 0.08f)));

        int x0 = r.x - margin - pad_x;
        int y0 = r.y - margin - pad_y;
        int x1 = r.x + r.width + margin - pad_x;
        int y1 = r.y + r.height + margin - pad_y;

        x0 = static_cast<int>(std::floor(x0 / scale));
        y0 = static_cast<int>(std::floor(y0 / scale));
        x1 = static_cast<int>(std::ceil(x1 / scale));
        y1 = static_cast<int>(std::ceil(y1 / scale));

        x0 = std::max(0, std::min(x0, orig_w - 1));
        y0 = std::max(0, std::min(y0, orig_h - 1));
        x1 = std::max(0, std::min(x1, orig_w));
        y1 = std::max(0, std::min(y1, orig_h));

        if (x1 <= x0 || y1 <= y0) {
            continue;
        }

        cv::Rect orig_rect(x0, y0, x1 - x0, y1 - y0);
        if (orig_rect.width < 3 || orig_rect.height < 3) {
            continue;
        }

        boxes.push_back({orig_rect, score});
    }

    std::sort(boxes.begin(), boxes.end(), [](const TextBox& a, const TextBox& b) {
        if (std::abs(a.rect.y - b.rect.y) > 10) {
            return a.rect.y < b.rect.y;
        }
        return a.rect.x < b.rect.x;
    });

    return boxes;
}

static cv::Mat prepare_rec_image(const cv::Mat& bgr_crop, int rec_w, int rec_h) {
    cv::Mat rgb;
    cv::cvtColor(bgr_crop, rgb, cv::COLOR_BGR2RGB);

    int src_w = rgb.cols;
    int src_h = rgb.rows;

    float ratio = static_cast<float>(rec_h) / src_h;
    int new_w = static_cast<int>(std::round(src_w * ratio));
    new_w = std::max(1, std::min(new_w, rec_w));

    cv::Mat resized;
    cv::resize(rgb, resized, cv::Size(new_w, rec_h), 0, 0, cv::INTER_LINEAR);

    cv::Mat padded(rec_h, rec_w, CV_8UC3, cv::Scalar(0, 0, 0));
    resized.copyTo(padded(cv::Rect(0, 0, new_w, rec_h)));

    return padded;
}

static std::string ctc_decode(const std::vector<float>& logits,
                              int time_steps,
                              int class_num,
                              const std::vector<std::string>& dict) {
    std::string text;
    int prev_idx = -1;

    for (int t = 0; t < time_steps; ++t) {
        const float* row = logits.data() + t * class_num;

        int max_idx = 0;
        float max_score = row[0];

        for (int c = 1; c < class_num; ++c) {
            if (row[c] > max_score) {
                max_score = row[c];
                max_idx = c;
            }
        }

        if (max_idx != 0 && max_idx != prev_idx) {
            if (max_idx >= 1 && max_idx <= static_cast<int>(dict.size())) {
                text += dict[max_idx - 1];
            } else if (max_idx == static_cast<int>(dict.size()) + 1) {
                text += " ";
            }
        }

        prev_idx = max_idx;
    }

    return text;
}

static void write_result(const std::string& path,
                         const std::vector<TextBox>& boxes,
                         const std::vector<std::string>& texts) {
    std::ofstream ofs(path);
    if (!ofs) {
        throw std::runtime_error("failed to write result: " + path);
    }

    for (size_t i = 0; i < boxes.size(); ++i) {
        const cv::Rect& r = boxes[i].rect;
        ofs << i
            << "\t" << r.x
            << "\t" << r.y
            << "\t" << r.width
            << "\t" << r.height
            << "\t" << boxes[i].score
            << "\t" << texts[i]
            << "\n";
    }
}

int main(int argc, char** argv) {
    if (argc < 6) {
        printf("Usage:\n");
        printf("  %s det.rknn rec_no_softmax.rknn dict.txt input.jpg output_dir\n", argv[0]);
        return -1;
    }

    const std::string det_model_path = argv[1];
    const std::string rec_model_path = argv[2];
    const std::string dict_path = argv[3];
    const std::string image_path = argv[4];
    const std::string output_dir = argv[5];

    try {
        cv::Mat image = cv::imread(image_path, cv::IMREAD_COLOR);
        if (image.empty()) {
            throw std::runtime_error("failed to read image: " + image_path);
        }

        std::vector<std::string> dict = load_dict(dict_path);
        printf("dict size: %zu\n", dict.size());

        RknnModel det_model(det_model_path);
        RknnModel rec_model(rec_model_path);

        print_tensor_shape("det input", det_model.input_attr());
        print_tensor_shape("det output", det_model.output_attr());
        print_tensor_shape("rec input", rec_model.input_attr());
        print_tensor_shape("rec output", rec_model.output_attr());

        const int det_w = 736;
        const int det_h = 736;

        float det_scale = 1.0f;
        int pad_x = 0;
        int pad_y = 0;

        cv::Mat det_bgr = resize_with_padding(image, det_w, det_h, det_scale, pad_x, pad_y);

        cv::Mat det_rgb;
        cv::cvtColor(det_bgr, det_rgb, cv::COLOR_BGR2RGB);

        std::vector<float> det_out = det_model.run_uint8_nhwc(
            det_rgb.data,
            det_w * det_h * 3
        );

        std::vector<TextBox> boxes = postprocess_det(
            det_out,
            det_w,
            det_h,
            image.cols,
            image.rows,
            det_scale,
            pad_x,
            pad_y
        );

        printf("detected boxes: %zu\n", boxes.size());

        cv::Mat vis = image.clone();
        std::vector<std::string> texts;

        const int rec_w = 320;
        const int rec_h = 48;
        const int rec_t = 40;
        const int rec_c = 18710;

        std::string crop_dir = output_dir + "/crops";
        std::string mkdir_cmd = "mkdir -p " + crop_dir;
        system(mkdir_cmd.c_str());

        for (size_t i = 0; i < boxes.size(); ++i) {
            cv::Rect r = boxes[i].rect & cv::Rect(0, 0, image.cols, image.rows);
            cv::Mat crop = image(r).clone();

            char crop_path[512];
            snprintf(crop_path, sizeof(crop_path), "%s/crop_%03zu.jpg", crop_dir.c_str(), i);
            cv::imwrite(crop_path, crop);

            cv::Mat rec_input = prepare_rec_image(crop, rec_w, rec_h);

            std::vector<float> rec_out = rec_model.run_uint8_nhwc(
                rec_input.data,
                rec_w * rec_h * 3
            );

            std::string text = ctc_decode(rec_out, rec_t, rec_c, dict);
            texts.push_back(text);

            cv::rectangle(vis, r, cv::Scalar(0, 255, 0), 2);
            cv::putText(vis,
                        std::to_string(i),
                        cv::Point(r.x, std::max(0, r.y - 3)),
                        cv::FONT_HERSHEY_SIMPLEX,
                        0.6,
                        cv::Scalar(0, 0, 255),
                        2);

            printf("[%zu] score=%.4f box=(%d,%d,%d,%d) text=%s\n",
                   i,
                   boxes[i].score,
                   r.x,
                   r.y,
                   r.width,
                   r.height,
                   text.c_str());
        }

        std::string result_txt = output_dir + "/ocr_result.txt";
        std::string result_img = output_dir + "/ocr_result.jpg";

        write_result(result_txt, boxes, texts);
        cv::imwrite(result_img, vis);

        printf("saved: %s\n", result_txt.c_str());
        printf("saved: %s\n", result_img.c_str());

        return 0;

    } catch (const std::exception& e) {
        printf("exception: %s\n", e.what());
        return -1;
    }
}
