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

#include <cinttypes>
#include <cmath>
#include <csignal>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <getopt.h>
#include <string>
#include <sys/stat.h>
#include <thread>
#include <unistd.h>
#include <vector>

static volatile bool g_interrupted = false;

static void noop_log(enum ggml_log_level level, const char * text, void * user_data) {
  (void)level; (void)text; (void)user_data;
}

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

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

static void batch_add(llama_batch &batch, const llama_token id,
                      const llama_pos pos,
                      const std::vector<llama_seq_id> &seq_ids,
                      const bool logits) {
  batch.token[batch.n_tokens] = id;
  batch.pos[batch.n_tokens] = pos;
  batch.n_seq_id[batch.n_tokens] = static_cast<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 llama_vocab *vocab,
                                  const llama_token token) {
  std::string piece;
  piece.resize(256);
  int32_t n = llama_token_to_piece(
      vocab, token, &piece[0], static_cast<int32_t>(piece.size()), 0, false);
  if (n < 0) {
    piece.resize(-n);
    n = llama_token_to_piece(vocab, token, &piece[0],
                             static_cast<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 (const 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;
  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 std::string capture_screenshot(bool quiet) {
  const char *tmpdir = getenv("TMPDIR");
  if (!tmpdir)
    tmpdir = "/tmp";
  std::string path = std::string(tmpdir) + "/ocr_screenshot.png";

  if (!quiet) fprintf(stderr, "Select a screen region to OCR...\n");

  auto check = [&]() -> bool {
    struct stat st {};
    return stat(path.c_str(), &st) == 0 && st.st_size > 0;
  };

  // 1. gnome-screenshot (blocks until selection complete)
  unlink(path.c_str());
  if (system(("gnome-screenshot --area -f " + path + " >/dev/null 2>&1")
                 .c_str()) == 0 &&
      check())
    return path;

  // 2. spectacle (KDE, blocks)
  unlink(path.c_str());
  if (system(
          ("spectacle --region -b -o " + path + " >/dev/null 2>&1").c_str()) ==
          0 &&
      check())
    return path;

  // 3. flameshot (writes raw PNG to stdout, blocks)
  unlink(path.c_str());
  if (system(("flameshot gui -r > " + path + " 2>/dev/null").c_str()) == 0 &&
      check())
    return path;

  // 4. maim (lightweight X11, blocks)
  unlink(path.c_str());
  if (system(("maim -s " + path + " 2>/dev/null").c_str()) == 0 && check())
    return path;

  // 5. slurp + grim (Sway/Wayland, blocks)
  unlink(path.c_str());
  if (system(("slurp | grim -g - " + path + " 2>/dev/null").c_str()) == 0 &&
      check())
    return path;

  return "";
}

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"
      "      --chat-template <name>  chat template (default: auto from model)\n"
      "      --screenshot         interactively select screen region to OCR\n"
      "  -q, --quiet             suppress info logs, show only OCR result\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"
      "  %s -m model.gguf --mmproj mmproj.gguf --screenshot   # select "
      "region\n",
      prog, 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";
  std::string chat_template;
  bool screenshot_mode = false;
  int n_threads = static_cast<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;
  bool quiet = 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 (true) {
    static 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'},
        {"chat-template", required_argument, nullptr, 4},
        {"screenshot", no_argument, nullptr, 5},
        {"quiet", no_argument, nullptr, 6},
        {"help", no_argument, nullptr, 'h'},
        {nullptr, 0, nullptr, 0}};
    int idx = 0;
    int c = getopt_long(argc, argv, "m:p:t:c:s:qh", 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 4:
      chat_template = optarg;
      break;
    case 5:
      screenshot_mode = true;
      break;
    case 6:
    case 'q':
      quiet = true;
      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 = static_cast<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 (screenshot_mode) {
    image_path = capture_screenshot(quiet);
    if (image_path.empty()) {
      fprintf(stderr,
              "ERROR: no screenshot tool found. Install gnome-screenshot, "
              "flameshot, spectacle, maim, or slurp+grim.\n");
      return 1;
    }
  } else if (optind < argc) {
    image_path = argv[optind];
  }

  signal(SIGINT, sigint_handler);

  int64_t t_start_us = ggml_time_us();

  llama_backend_init();

  if (quiet) {
    llama_log_set(noop_log, nullptr);
    mtmd_helper_log_set(noop_log, nullptr);
  }

  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 = static_cast<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 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;
  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());
      if (screenshot_mode)
        unlink(image_path.c_str());
      mtmd_free(ctx_vision);
      llama_free(lctx);
      llama_model_free(model);
      return 1;
    }
  } else {
    if (!quiet) 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;
    }
  }

  // Build user message with image marker placeholder
  std::string user_content = std::string(mtmd_default_marker()) + prompt;

  // Apply chat template for any model
  std::string full_prompt;
  const char *tmpl = chat_template.empty()
                         ? llama_model_chat_template(model, nullptr)
                         : chat_template.c_str();
  if (tmpl) {
    llama_chat_message msg[1] = {{"user", user_content.c_str()}};
    char buf[8192];
    int32_t n = llama_chat_apply_template(tmpl, msg, 1, true, buf, sizeof(buf));
    if (n > 0) {
      full_prompt = std::string(buf, n);
    }
  }

  // Fallback: use chatml format for models without recognizable template
  if (full_prompt.empty()) {
    full_prompt = "<|im_start|>user\n" + user_content +
                  "<|im_end|>\n<|im_start|>assistant\n";
  }

  // Template output already includes role markers and special tokens (e.g.
  // [gMASK]<sop>) So we don't add BOS separately — add_special=false
  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;
  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);
    if (screenshot_mode)
      unlink(image_path.c_str());
    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);
    if (screenshot_mode)
      unlink(image_path.c_str());
    mtmd_free(ctx_vision);
    llama_free(lctx);
    llama_model_free(model);
    return 1;
  }

  mtmd_input_chunks_free(chunks);
  mtmd_bitmap_free(bmp);
  if (screenshot_mode)
    unlink(image_path.c_str());

  llama_sampler *smpl;
  if (temp <= 0.0f) {
    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 {
    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));
  }

  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",
           static_cast<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;
}