#include <DwarfInterpret.hpp>

#include <fstream>
#include <iostream>
#include <iomanip>
#include <vector>
#include <set>

#include <fileno.hpp>
#include <dwarfpp/lib.hpp>
#include <dwarfpp/regs.hpp>
#include <dwarfpp/frame.hpp>
#include <dwarfpp/attr.hpp>
#include <dwarfpp/frame.hpp>
#include <dwarfpp/root.hpp>
#include <gelf.h>

#include <ucontext.h>

#include "register_map.hpp"

using namespace std;
using namespace dwarf;


std::unique_ptr<DwarfInterpret> DwarfInterpret::instance(nullptr);

DwarfInterpret::DwarfInterpret(const std::string& elf_path)
    : root_die(fileno(std::ifstream(elf_path)))
{
	//std::ifstream file_in(elf_path);
	//root_die = dwarf::core::root_die(fileno(file_in));
    // FIXME this ^ deserves some checks. But this looks tedious.

	GElf_Ehdr ehdr;
	GElf_Ehdr *ret = gelf_getehdr(root_die.get_elf(), &ehdr);
	assert(ret != 0);
	elf_machine = ehdr.e_machine;
}

DwarfInterpret& DwarfInterpret::acquire() {
    if(DwarfInterpret::instance == nullptr)
        throw NotInstanciated();
    return *(DwarfInterpret::instance);
}

DwarfInterpret& DwarfInterpret::instanciate(const std::string& elf_path) {
    if(DwarfInterpret::instance != nullptr)
        throw AlreadyInstanciated();
    DwarfInterpret::instance = std::unique_ptr<DwarfInterpret>(
            new DwarfInterpret(elf_path));
    return *(DwarfInterpret::instance);
}

DwarfInterpret::reg_content_t DwarfInterpret::interpret_dw_register(
        const DwarfInterpret::DwarfRow& row,
        const DwarfInterpret::DwarfRegister& reg
        ) const
{
    switch(reg.k) {
        case core::FrameSection::register_def::INDETERMINATE:
        case core::FrameSection::register_def::UNDEFINED:
            throw ValuelessRegister();

        case core::FrameSection::register_def::REGISTER: {
            ostringstream exprs;
            int dwarf_regnum = reg.register_plus_offset_r().first;
            assert(dwarf_regnum != DW_FRAME_CFA_COL3);
            int cpu_regnum = REG_DWARF_TO_UCONTEXT[dwarf_regnum];

            reg_content_t* cpu_content =
                (reg_content_t*) get_cpu_register(cpu_regnum);
            reg_content_t* addr =
                cpu_content + reg.register_plus_offset_r().second;

            return *addr;
        } break;

        case core::FrameSection::register_def::SAVED_AT_OFFSET_FROM_CFA: {
            reg_content_t* cfa_loc =
                (reg_content_t*) interpret_dw_register(row, DW_FRAME_CFA_COL3);
            reg_content_t* addr = cfa_loc;
            return *addr;
        } break;

        case core::FrameSection::register_def::SAVED_AT_EXPR:
            throw NotImplemented();
            break;
        case core::FrameSection::register_def::VAL_IS_OFFSET_FROM_CFA:
        case core::FrameSection::register_def::VAL_OF_EXPR:
        default:
            throw NotImplemented();
            break;
    }
}

DwarfInterpret::reg_content_t DwarfInterpret::interpret_dw_register(
        const DwarfInterpret::DwarfRow& row,
        int reg_id
        ) const
{
    return interpret_dw_register(row, get_column(row, reg_id));
}

uintptr_t DwarfInterpret::get_return_address(uintptr_t cur_pc) const {
    const core::Cie& cie = *cie_at(cur_pc);
    const DwarfRow& row = dwarf_row_at(cur_pc);

    return interpret_dw_register(row, cie.get_return_address_register_rule());
}

uintptr_t DwarfInterpret::get_self_return_address() const {
    // Aaaaand now we have to get_return_address thrice.
    return get_return_address(get_caller_pc());
}

uintptr_t DwarfInterpret::get_current_pc() {
    // Assumes the PC is stored in REG_RIP **AND** the PC we want is this
    // function's return address

    DwarfInterpret& dw = DwarfInterpret::acquire();
    reg_content_t pc_here = dw.get_cpu_register(REG_RIP);
    fprintf(stderr, "PC=%p <inside>\n", pc_here);
    fprintf(stderr, "actual PC=%p <inside>\n", __builtin_return_address(0));
    return dw.get_return_address(pc_here);
}

uintptr_t DwarfInterpret::get_caller_pc() const {
    // We assume we want the PC of the caller of the calling function. This
    // means we have to unwind twice. `get_current_pc` unwinds once.

    return get_return_address(DwarfInterpret::get_current_pc());
}

template<typename Key>
static typename std::set<std::pair<int, Key> >::const_iterator find_column(
        const std::set<std::pair<int, Key> >& set,
        int key)
{
    if(set.empty())
        return set.end();
    Key some_key_val = set.begin()->second;
    auto it = set.upper_bound(make_pair(key - 1, some_key_val));
    if(it == set.end())
        return set.end();
    while(it != set.end() && it->first != key)
        ++it;
    if(it == set.end() || it->first != key)
        return set.end();
    return it;
}

DwarfInterpret::DwarfRegister DwarfInterpret::get_column(
        const DwarfInterpret::DwarfRow& row,
        int column
        ) const
{
    auto it = find_column<DwarfRegister>(row, column);
    if(it == row.end())
        throw std::out_of_range("No such column");
    return it->second;
}

DwarfInterpret::reg_content_t DwarfInterpret::get_cpu_register(
        int reg_id) const
{
    ucontext_t context;
    if(getcontext(&context) != 0)
        throw FailedGetContext();

    // Let the user deal with reg_id correctness
    return context.uc_mcontext.gregs[reg_id];
}

const core::FrameSection::fde_iterator DwarfInterpret::fde_at(
        uintptr_t pc
        ) const
{
    const core::FrameSection& fs = root_die.get_frame_section();
    const auto& fde_it = fs.find_fde_for_pc(pc);
    if(fde_it == fs.fde_end())
        throw NotFound(std::string("FDE at pc=") + std::to_string(pc));

    return fde_it;
}

const core::FrameSection::cie_iterator DwarfInterpret::cie_at(
        uintptr_t pc
        ) const
{
    const core::Fde& fde = *fde_at(pc); // let NotFound escape if it occurs
    const auto& cie_it = fde.find_cie();
    if(cie_it == root_die.get_frame_section().cie_end()) // Should not happen…?
        throw NotFound(std::string("CIE matching FDE at pc=")
                + std::to_string(pc));

    return cie_it;
}

const DwarfInterpret::DwarfRow& DwarfInterpret::dwarf_row_at(
        uintptr_t pc
        ) const
{
    const core::Fde& fde = *fde_at(pc);
    auto decoded = fde.decode();
    const auto& row_it = decoded.rows.find(pc);
    if(row_it == decoded.rows.end())
        throw NotFound(std::string("Dwarf row in this FDE at pc=")
                + std::to_string(pc));

    return row_it->second;
}