#include "tracer.hpp"
#include <cerrno>
#include <cstring>
#include <filesystem>
#include <format>
#include <fstream>
#include <iostream>
#include <sstream>
#include <sys/ptrace.h>
#include <sys/user.h>
#include <sys/wait.h>
#include <unistd.h>

// Hardware watchpoint condition codes
#define DR7_BREAK_ON_WRITE 0x1
#define DR7_BREAK_ON_RW 0x3

// Hardware watchpoint size codes
#define DR7_LEN_1 0x0
#define DR7_LEN_2 0x1
#define DR7_LEN_4 0x3
#define DR7_LEN_8 0x2

void Tracer::setup_hardware_watchpoint(pid_t pid, uint64_t address, size_t size) {
  // Use DR0 for our watchpoint
  ptrace(PTRACE_POKEUSER, pid, offsetof(struct user, u_debugreg[0]), address);

  // Configure DR7 for DR0
  uint64_t dr7 = 0;

  // Set local enable bit for DR0 (bit 0)
  dr7 |= (1 << 0);

  // Set condition to break on read/write (bits 16-17 for DR0)
  dr7 |= (DR7_BREAK_ON_RW << 16);

  // Set length based on size (bits 18-19 for DR0)
  uint64_t len_code;
  if (size == 1) len_code = DR7_LEN_1;
  else if (size == 2) len_code = DR7_LEN_2;
  else if (size == 4) len_code = DR7_LEN_4;
  else len_code = DR7_LEN_8;

  dr7 |= (len_code << 18);

  // Also set GE (bit 9) for exact breakpoint detection
  dr7 |= (1 << 9);

  // Write DR7
  ptrace(PTRACE_POKEUSER, pid, offsetof(struct user, u_debugreg[7]), dr7);
}

int64_t Tracer::read_memory_value(pid_t pid, uint64_t address, size_t size) {
  int64_t value = 0;

  if (size == 1) {
    value = ptrace(PTRACE_PEEKDATA, pid, address, 0) & 0xFF;
  } else if (size == 2) {
    value = ptrace(PTRACE_PEEKDATA, pid, address, 0) & 0xFFFF;
  } else if (size == 4) {
    value = ptrace(PTRACE_PEEKDATA, pid, address, 0) & 0xFFFFFFFF;
  } else if (size == 8) {
    value = ptrace(PTRACE_PEEKDATA, pid, address, 0);
  }

  return value;
}

void Tracer::handle_watchpoint_hit(pid_t pid, uint64_t address, size_t size) {
  int64_t current_value = read_memory_value(pid, address, size);

  if (!prev_value_.has_value()) {
    // First access - initialize prev_value
    prev_value_ = current_value;
    std::cout << var_name_ << "\tread\t" << current_value << std::endl;
  } else if (prev_value_.value() != current_value) {
    // Value changed - this was a write
    std::cout << var_name_ << "\twrite\t" << prev_value_.value()
              << "->" << current_value << std::endl;
    prev_value_ = current_value;
  } else {
    // Value unchanged - this was a read
    std::cout << var_name_ << "\tread\t" << current_value << std::endl;
  }
}

uint64_t Tracer::get_load_address(pid_t pid, const std::string &exec_path) {
  std::string maps_path = std::format("/proc/{}/maps", pid);
  std::ifstream maps_file(maps_path);

  if (!maps_file.is_open()) {
    std::cerr << "Error: Cannot open " << maps_path << std::endl;
    return 0;
  }

  std::filesystem::path exec_name = std::filesystem::path(exec_path).filename();
  std::string line;

  while (std::getline(maps_file, line)) {
    // Look for executable mapping containing the binary name
    if (line.find(exec_name.string()) != std::string::npos &&
        line.find(" r-xp ") != std::string::npos) {
      // Parse the address range
      std::istringstream iss(line);
      std::string addr_range;
      iss >> addr_range;

      // Extract start address
      size_t dash_pos = addr_range.find('-');
      if (dash_pos != std::string::npos) {
        std::string start_addr_str = addr_range.substr(0, dash_pos);
        return std::stoull(start_addr_str, nullptr, 16);
      }
    }
  }

  return 0;
}

void Tracer::watch_variable(const std::string &var_name, const std::string &exec_path,
                            uint64_t address, size_t size,
                            const std::vector<std::string> &args) {
  var_name_ = var_name;
  prev_value_.reset();

  pid_t pid = fork();

  if (pid == 0) {
    // Child process: execute the target binary
    ptrace(PTRACE_TRACEME, 0, 0, 0);

    // Build argv array for execv
    std::vector<char*> argv_vec;
    argv_vec.push_back(const_cast<char*>(exec_path.c_str()));
    for (const auto& arg : args) {
      argv_vec.push_back(const_cast<char*>(arg.c_str()));
    }
    argv_vec.push_back(nullptr);

    execv(exec_path.c_str(), argv_vec.data());
    std::cerr << "Error: Failed to execute " << exec_path << std::endl;
    exit(1);
  } else if (pid > 0) {
    // Parent process: set up watchpoint and monitor
    int status;
    waitpid(pid, &status, 0);

    if (!WIFSTOPPED(status)) {
      std::cerr << "Error: Child process not stopped after PTRACE_TRACEME" << std::endl;
      return;
    }

    // The process is stopped after exec due to PTRACE_TRACEME
    // Get the actual load address (for PIE binaries)
    uint64_t load_addr = get_load_address(pid, exec_path);
    uint64_t runtime_addr = load_addr + address;

    std::cout << std::format("Load address: 0x{:x}", load_addr) << std::endl;
    std::cout << std::format("Runtime address: 0x{:x}", runtime_addr) << std::endl;

    // Set up hardware watchpoint
    setup_hardware_watchpoint(pid, runtime_addr, size);

    // Continue execution and monitor for watchpoint hits
    ptrace(PTRACE_CONT, pid, 0, 0);

    while (true) {
      waitpid(pid, &status, 0);

      if (WIFEXITED(status)) {
        std::cout << "Process exited with status " << WEXITSTATUS(status)
                  << std::endl;
        break;
      }

      if (WIFSTOPPED(status)) {
        int sig = WSTOPSIG(status);

        if (sig == SIGTRAP) {
          // Check if it's a watchpoint hit
          uint64_t dr6 = ptrace(PTRACE_PEEKUSER, pid,
                               offsetof(struct user, u_debugreg[6]), 0);

          if (dr6 & 0x1) {  // DR0 triggered
            handle_watchpoint_hit(pid, runtime_addr, size);

            // Clear DR6
            ptrace(PTRACE_POKEUSER, pid,
                   offsetof(struct user, u_debugreg[6]), 0);
          }

          ptrace(PTRACE_CONT, pid, 0, 0);
        } else {
          // Forward other signals
          ptrace(PTRACE_CONT, pid, 0, sig);
        }
      }
    }
  } else {
    std::cerr << "Error: fork() failed" << std::endl;
  }
}