#include "llama.h"
#include "mtmd.h"
#include "mtmd-helper.h"
#include "ggml.h"

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cmath>
#include <string>
#include <vector>
#include <thread>
#include <unistd.h>
#include <getopt.h>
#include <signal.h>
#include <inttypes.h>

static volatile bool g_interrupted = false;

static void sigint_handler(int signo) {
    (void)signo;
    g_interrupted = true;
}

static void batch_clear(struct llama_batch & batch) {
    batch.n_tokens = 0;
}

static void batch_add(struct llama_batch & batch, llama_token id, llama_pos pos, const std::vector<llama_seq_id> & seq_ids, bool logits) {
    batch.token   [batch.n_tokens] = id;
    batch.pos     [batch.n_tokens] = pos;
    batch.n_seq_id[batch.n_tokens] = (int32_t)seq_ids.size();
    for (size_t i = 0; i < seq_ids.size(); ++i) {
        batch.seq_id[batch.n_tokens][i] = seq_ids[i];
    }
    batch.logits  [batch.n_tokens] = logits ? 1 : 0;
    batch.n_tokens++;
}

static std::string token_to_piece(const struct llama_vocab * vocab, llama_token token) {
    std::string piece;
    piece.resize(256);
    int32_t n = llama_token_to_piece(vocab, token, &piece[0], (int32_t)piece.size(), 0, false);
    if (n < 0) {
        piece.resize(-n);
        n = llama_token_to_piece(vocab, token, &piece[0], (int32_t)piece.size(), 0, false);
    }
    piece.resize(std::max(0, n));
    return piece;
}

static std::string json_escape(const std::string & s) {
    std::string out;
    out.reserve(s.size() + 2);
    for (char c : s) {
        switch (c) {
            case '"':  out += "\\\""; break;
            case '\\': out += "\\\\"; break;
            case '\n': out += "\\n";  break;
            case '\r': out += "\\r";  break;
            case '\t': out += "\\t";  break;
            default:   out += c;      break;
        }
    }
    return out;
}

static std::vector<unsigned char> read_clipboard() {
    const char * cmd = nullptr;
    if (getenv("WAYLAND_DISPLAY")) {
        cmd = "wl-paste 2>/dev/null";
    } else {
        cmd = "xclip -selection clipboard -t image/png -o 2>/dev/null";
    }
    std::vector<unsigned char> data;
    FILE * pipe = popen(cmd, "re");
    if (!pipe) return data;
    unsigned char buf[65536];
    size_t n;
    while ((n = fread(buf, 1, sizeof(buf), pipe)) > 0) {
        data.insert(data.end(), buf, buf + n);
    }
    int st = pclose(pipe);
    if (st != 0 || data.empty()) {
        data.clear();
        // fallback for X11: try without explicit target
        if (!getenv("WAYLAND_DISPLAY")) {
            pipe = popen("xclip -selection clipboard -o 2>/dev/null", "re");
            if (pipe) {
                while ((n = fread(buf, 1, sizeof(buf), pipe)) > 0)
                    data.insert(data.end(), buf, buf + n);
                pclose(pipe);
            }
        }
    }
    return data;
}

static void print_usage(const char * prog) {
    fprintf(stderr,
        "Usage: %s -m <model> --mmproj <mmproj> [options] [<image_path>]\n"
        "\n"
        "Options:\n"
        "  -m, --model  <path>    model file path (GGUF)\n"
        "      --mmproj,--mm <path>  mmproj file path\n"
        "  -p           <text>    OCR prompt (default: OCR)\n"
        "  -t           <n>       number of threads (default: cpu cores)\n"
        "      --ngl    <n>       GPU layers (-1 = all, default: -1)\n"
        "  -c           <n>       context size (default: 8192)\n"
        "      --temp   <f>       sampling temperature (0 = greedy, default: 0)\n"
        "  -s           <seed>    random seed\n"
        "      --json              output in JSON format\n"
        "  -h                      show this help\n"
        "\n"
        "If <image_path> is omitted, reads image from clipboard.\n"
        "Requires wl-paste (Wayland) or xclip (X11) for clipboard support.\n"
        "\n"
        "Examples:\n"
        "  %s -m model.gguf --mmproj mmproj.gguf -p \"OCR\" image.png\n"
        "  %s -m model.gguf --mmproj mmproj.gguf   # use clipboard image\n",
        prog, prog, prog);
}

int main(int argc, char ** argv) {
    std::string model_path;
    std::string mmproj_path;
    std::string image_path;
    std::string prompt = "OCR";
    int n_threads = (int)std::thread::hardware_concurrency();
    int n_gpu_layers = -1;
    int n_ctx = 8192;
    float temp = 0.0f;
    uint32_t seed = LLAMA_DEFAULT_SEED;
    bool json_output = false;

    // pre-process -mm to --mmproj so getopt handles it correctly
    for (int i = 1; i < argc; i++) {
        if (strcmp(argv[i], "-mm") == 0) {
            argv[i] = (char *)"--mmproj";
        }
    }

    while (1) {
        static struct option long_opts[] = {
            {"temp",   required_argument, nullptr, 0},
            {"json",   no_argument,       nullptr, 1},
            {"mmproj", required_argument, nullptr, 2},
            {"mm",     required_argument, nullptr, 2},
            {"ngl",    required_argument, nullptr, 3},
            {"model",  required_argument, nullptr, 'm'},
            {"help",   no_argument,       nullptr, 'h'},
            {nullptr, 0, nullptr, 0}
        };
        int idx = 0;
        int c = getopt_long(argc, argv, "m:p:t:c:s:h", long_opts, &idx);
        if (c == -1) break;
        switch (c) {
            case 0: temp = std::stof(optarg); break;
            case 1: json_output = true; break;
            case 2: mmproj_path = optarg; break;
            case 3: n_gpu_layers = std::stoi(optarg); break;
            case 'm': model_path = optarg; break;
            case 'p': prompt = optarg; break;
            case 't': n_threads = std::stoi(optarg); break;
            case 'c': n_ctx = std::stoi(optarg); break;
            case 's': seed = (uint32_t)std::stoul(optarg); break;
            case 'h': print_usage(argv[0]); return 0;
            default:  print_usage(argv[0]); return 1;
        }
    }

    if (model_path.empty() || mmproj_path.empty()) {
        fprintf(stderr, "ERROR: -m/--model and --mmproj/--mm are required\n");
        return 1;
    }

    if (optind < argc) {
        image_path = argv[optind];
    }

    signal(SIGINT, sigint_handler);

    int64_t t_start_us = ggml_time_us();

    llama_backend_init();

    llama_model_params mparams = llama_model_default_params();
    mparams.n_gpu_layers = n_gpu_layers;

    llama_model * model = llama_model_load_from_file(model_path.c_str(), mparams);
    if (!model) {
        fprintf(stderr, "ERROR: failed to load model from %s\n", model_path.c_str());
        return 1;
    }

    llama_context_params cparams = llama_context_default_params();
    cparams.n_ctx     = (uint32_t)n_ctx;
    cparams.n_batch   = 512;
    cparams.n_ubatch  = 512;

    llama_context * lctx = llama_init_from_model(model, cparams);
    if (!lctx) {
        fprintf(stderr, "ERROR: failed to create context\n");
        llama_model_free(model);
        return 1;
    }

    llama_set_n_threads(lctx, n_threads, n_threads);

    const struct llama_vocab * vocab = llama_model_get_vocab(model);

    mtmd_context_params mtmd_params = mtmd_context_params_default();
    mtmd_params.use_gpu   = (n_gpu_layers != 0);
    mtmd_params.n_threads = n_threads;

    mtmd_context * ctx_vision = mtmd_init_from_file(mmproj_path.c_str(), model, mtmd_params);
    if (!ctx_vision) {
        fprintf(stderr, "ERROR: failed to load mmproj from %s\n", mmproj_path.c_str());
        llama_free(lctx);
        llama_model_free(model);
        return 1;
    }

    int64_t t_loaded_us = ggml_time_us();

    mtmd_bitmap * bmp = nullptr;
    if (!image_path.empty()) {
        bmp = mtmd_helper_bitmap_init_from_file(ctx_vision, image_path.c_str());
        if (!bmp) {
            fprintf(stderr, "ERROR: failed to load image from %s\n", image_path.c_str());
            mtmd_free(ctx_vision);
            llama_free(lctx);
            llama_model_free(model);
            return 1;
        }
    } else {
        fprintf(stderr, "Reading image from clipboard...\n");
        auto clip = read_clipboard();
        if (clip.empty()) {
            fprintf(stderr, "ERROR: clipboard is empty or no image found.\n"
                            "Make sure wl-paste (Wayland) or xclip (X11) is installed.\n");
            mtmd_free(ctx_vision);
            llama_free(lctx);
            llama_model_free(model);
            return 1;
        }
        bmp = mtmd_helper_bitmap_init_from_buf(ctx_vision, clip.data(), clip.size());
        if (!bmp) {
            fprintf(stderr, "ERROR: failed to decode clipboard image.\n");
            mtmd_free(ctx_vision);
            llama_free(lctx);
            llama_model_free(model);
            return 1;
        }
    }

    // Construct prompt using GLM-OCR chat template format
    // The model expects: [gMASK]<sop><|user|>\n<__media__>PROMPT\n<|assistant|>\n
    std::string full_prompt = "[gMASK]<sop><|user|>\n";
    full_prompt += mtmd_default_marker();
    full_prompt += prompt;
    full_prompt += "\n<|assistant|>\n";

    mtmd_input_chunks * chunks = mtmd_input_chunks_init();

    mtmd_input_text in_text;
    in_text.text          = full_prompt.c_str();
    in_text.add_special   = false; // template already includes [gMASK]/<sop>
    in_text.parse_special = true;

    const mtmd_bitmap * bitmaps[1] = { bmp };

    int32_t res = mtmd_tokenize(ctx_vision, chunks, &in_text, bitmaps, 1);
    if (res != 0) {
        fprintf(stderr, "ERROR: mtmd_tokenize failed (%d)\n", res);
        mtmd_input_chunks_free(chunks);
        mtmd_bitmap_free(bmp);
        mtmd_free(ctx_vision);
        llama_free(lctx);
        llama_model_free(model);
        return 1;
    }

    llama_pos n_past = 0;
    res = mtmd_helper_eval_chunks(ctx_vision, lctx, chunks, n_past, 0, 512, true, &n_past);
    if (res != 0) {
        fprintf(stderr, "ERROR: mtmd_helper_eval_chunks failed (%d)\n", res);
        mtmd_input_chunks_free(chunks);
        mtmd_bitmap_free(bmp);
        mtmd_free(ctx_vision);
        llama_free(lctx);
        llama_model_free(model);
        return 1;
    }

    mtmd_input_chunks_free(chunks);
    mtmd_bitmap_free(bmp);

    struct llama_sampler * smpl = nullptr;
    if (temp <= 0.0f) {
        struct llama_sampler_chain_params sparams = llama_sampler_chain_default_params();
        smpl = llama_sampler_chain_init(sparams);
        llama_sampler_chain_add(smpl, llama_sampler_init_greedy());
    } else {
        struct llama_sampler_chain_params sparams = llama_sampler_chain_default_params();
        smpl = llama_sampler_chain_init(sparams);
        llama_sampler_chain_add(smpl, llama_sampler_init_top_k(40));
        llama_sampler_chain_add(smpl, llama_sampler_init_top_p(0.9f, 1));
        llama_sampler_chain_add(smpl, llama_sampler_init_temp(temp));
        llama_sampler_chain_add(smpl, llama_sampler_init_dist(seed));
    }

    struct llama_batch batch = llama_batch_init(1, 0, 1);

    int64_t t_infer_start_us = ggml_time_us();
    std::string ocr_text;
    int n_generated = 0;

    for (int i = 0; i < 2048; i++) {
        if (g_interrupted) break;

        llama_token token_id = llama_sampler_sample(smpl, lctx, -1);
        llama_sampler_accept(smpl, token_id);

        if (llama_vocab_is_eog(vocab, token_id)) break;

        std::string piece = token_to_piece(vocab, token_id);
        ocr_text += piece;

        if (!json_output) {
            printf("%s", piece.c_str());
            fflush(stdout);
        }

        batch_clear(batch);
        batch_add(batch, token_id, n_past++, {0}, true);

        if (llama_decode(lctx, batch) != 0) {
            fprintf(stderr, "\nERROR: llama_decode failed\n");
            break;
        }

        n_generated++;
    }

    int64_t t_end_us = ggml_time_us();

    if (!json_output) {
        printf("\n");
    }

    llama_batch_free(batch);
    llama_sampler_free(smpl);
    mtmd_free(ctx_vision);
    llama_free(lctx);
    llama_model_free(model);
    llama_backend_free();

    if (json_output) {
        printf("{\n");
        printf("  \"text\": \"%s\",\n", json_escape(ocr_text).c_str());
        printf("  \"meta\": {\n");
        printf("    \"model\": \"%s\",\n", json_escape(model_path).c_str());
        printf("    \"mmproj\": \"%s\",\n", json_escape(mmproj_path).c_str());
        printf("    \"prompt\": \"%s\",\n", json_escape(prompt).c_str());
        printf("    \"n_prompt_tokens\": %" PRId64 ",\n", (int64_t)n_past);
        printf("    \"n_generated\": %d,\n", n_generated);
        printf("    \"timings_ms\": {\n");
        printf("      \"load\": %" PRId64 ",\n", (t_loaded_us - t_start_us) / 1000);
        printf("      \"inference\": %" PRId64 "\n", (t_end_us - t_infer_start_us) / 1000);
        printf("    }\n");
        printf("  }\n");
        printf("}\n");
    }

    return g_interrupted ? 130 : 0;
}