#include <dwarfinterpret/DwarfInterpret.hpp>

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

#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;


#define get_cpu_register(reg_id, dest) {\
    ucontext_t context; \
    if(getcontext(&context) != 0) \
        throw DwarfInterpret::FailedGetContext(); \
    (dest) = context.uc_mcontext.gregs[(reg_id)]; \
}



MemoryMap DwarfInterpret::memory_map;
map<int, unique_ptr<DwarfInterpret> > DwarfInterpret::instances;

DwarfInterpret::DwarfInterpret(const MemoryMap::MapEntry& memory_object)
    : map_entry(memory_object)
{
    if(memory_object.mem_region.begin >= memory_object.mem_region.end) {
        // We're building the anonymous interpret: don't do anything
        root_die = nullptr;
    }
    else {
        const string& path = memory_object.pathname;
        if(path[0] == '[') {
            throw InvalidElf(
                    string("This pathname is a special object: ") + path);
        }
        std::ifstream handle(path);
        if(!handle.good())
            throw InvalidElf(
                    string("Cannot read this ELF file: ") + path);

        root_die = unique_ptr<core::root_die>(
                new core::root_die(fileno(handle)));
        // TODO catch exceptions here ^^^^ and throw NoDebugData/InvalidElf

        pc_offset = memory_object.mem_region.begin - memory_object.offset;
    }
	//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() {
    auto find_it = DwarfInterpret::instances.find(-1);
    if(find_it == DwarfInterpret::instances.end()) {
        MemoryMap::MapEntry anon_map_entry;
        anon_map_entry.offset = 0;
        anon_map_entry.mem_region.begin = 0;
        anon_map_entry.mem_region.end = 0;
        DwarfInterpret* instance = new DwarfInterpret(anon_map_entry);
        DwarfInterpret::instances.insert(make_pair(-1,
                    unique_ptr<DwarfInterpret>(instance)));
        return *instance;
    }
    return *(find_it->second);
}

DwarfInterpret& DwarfInterpret::acquire(uintptr_t pc) {
    try {
        size_t entry_id = DwarfInterpret::memory_map.id_of_address(pc);
        auto find_it = DwarfInterpret::instances.find(entry_id);
        if(find_it == DwarfInterpret::instances.end()) {
            // Instanciate it
            MemoryMap::MapEntry map_entry =
                DwarfInterpret::memory_map[entry_id];
            DwarfInterpret* instance = new DwarfInterpret(map_entry);
            DwarfInterpret::instances.insert(make_pair(entry_id,
                        unique_ptr<DwarfInterpret>(instance)));
            return *instance;
        }

        // Already instanciated, return it
        return *(find_it->second);
    } catch(const std::out_of_range& e) {
        ostringstream exn_text;
        exn_text << "No ELF mapped section found for memory address 0x"
                 << hex << pc;
        throw(NoElfForPC(exn_text.str()));
    }
}

DwarfInterpret::reg_content_t DwarfInterpret::interpret_dw_register(
        const DwarfInterpret::DwarfRow& row,
        const DwarfInterpret::DwarfRegister& reg,
        const UnwindContext& ctx
        ) 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);

            reg_content_t cpu_content = 9;
            cerr << "@@ Interpreting register ";
            switch(dwarf_regnum) {
                case lib::DWARF_X86_64_RIP:
                    cerr << "RIP";
                    cpu_content = (reg_content_t) ctx.rip;
                    break;
                case lib::DWARF_X86_64_RBP:
                    cerr << "RBP";
                    cpu_content = (reg_content_t) ctx.rbp;
                    break;
                case lib::DWARF_X86_64_RSP:
                    cerr << "RSP";
                    cpu_content = (reg_content_t) ctx.rsp;
                    break;
                default:
                    throw ExoticRegister(to_string(dwarf_regnum));
            }
            reg_content_t addr =
                cpu_content + reg.register_plus_offset_r().second;

            cerr << " = " << hex
                 << cpu_content << " + offset = " << addr //<< " -> " << *addr
                 << dec << endl;


            return addr;
        } break;

        case core::FrameSection::register_def::SAVED_AT_OFFSET_FROM_CFA: {
            reg_content_t cfa_loc =
                interpret_dw_register(row, DW_FRAME_CFA_COL3, ctx);
            int cfa_offset = reg.saved_at_offset_from_cfa_r();
            reg_content_t addr = cfa_loc + cfa_offset;
            reg_content_t value = *((reg_content_t*) addr);
            cerr << "@@ Interpreting CFA offset: CFA is " << hex
                 << cfa_loc << " + offset " << dec << cfa_offset << hex
                 << " = " << addr
                 << " -> " << value
                 << dec << endl;
            return value;
        } 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 UnwindContext& ctx
        ) const
{
    return interpret_dw_register(row, get_column(row, reg_id), ctx);
}

DwarfInterpret::UnwindContext DwarfInterpret::get_current_unwind_context() {
    UnwindContext ctx;
    get_cpu_register(REG_RIP, ctx.rip);
    get_cpu_register(REG_RSP, ctx.rsp);
    get_cpu_register(REG_RBP, ctx.rbp);

    cerr << "CREATING CONTEXT. %rsp=0x" << hex
         << ctx.rsp
         << ", %rbp=0x" << ctx.rbp
         << ", %rip=0x" << ctx.rip
         << dec << endl;

    DwarfInterpret& dw = DwarfInterpret::acquire(ctx.rip);

    return dw.unwind_context(ctx);
}

DwarfInterpret::UnwindContext DwarfInterpret::unwind_context(
        const DwarfInterpret::UnwindContext& ctx
        )
{
    DwarfInterpret* responsible = get_responsible_instance(ctx.rip);
    if(responsible != nullptr)
        return responsible->unwind_context(ctx);

    DwarfRow cur_row = dwarf_row_at(ctx.rip);
    const core::Cie& cie = *cie_at(ctx.rip);
    UnwindContext new_context;
    cerr << "Obtaining previous context as reg "
         << cie.get_return_address_register_rule()
         << " at current IP = "
         << hex << ctx.rip << endl;
    new_context.rip = interpret_dw_register(
            cur_row,
            cie.get_return_address_register_rule(),
            ctx);
    cerr << "Yielding " << hex << new_context.rip << dec << endl;
    new_context.rbp = interpret_dw_register(
            cur_row,
            lib::DWARF_X86_64_RBP,
            ctx);

    new_context.rsp = interpret_dw_register(
            cur_row,
            DW_FRAME_CFA_COL3,
            ctx);

    return new_context;
}

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;
}

const core::FrameSection::fde_iterator DwarfInterpret::fde_at(
        uintptr_t pc
        ) const
{
    DwarfInterpret* responsible = get_responsible_instance(pc);
    if(responsible != nullptr)
        return responsible->fde_at(pc);

    cerr << "Extracting FDE with ";
    map_entry.dump(cerr);

    const core::FrameSection& fs = root_die->get_frame_section();
    uintptr_t translated_pc = pc - pc_offset;
    const auto& fde_it = fs.find_fde_for_pc(translated_pc);
    if(fde_it == fs.fde_end()) {
        ostringstream exn_text;
        exn_text << "FDE at pc=0x" << hex << pc
                 << " (0x" << translated_pc << " in ELF)";
        throw NotFound(exn_text.str());
    }

    return fde_it;
}

const core::FrameSection::cie_iterator DwarfInterpret::cie_at(
        uintptr_t pc
        ) const
{
    DwarfInterpret* responsible = get_responsible_instance(pc);
    if(responsible != nullptr)
        return responsible->cie_at(pc);

    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;
}

DwarfInterpret::DwarfRow DwarfInterpret::dwarf_row_at(
        uintptr_t pc
        ) const
{
    DwarfInterpret* responsible = get_responsible_instance(pc);
    if(responsible != nullptr)
        return responsible->dwarf_row_at(pc);

    const core::Fde& fde = *fde_at(pc);
    uintptr_t translated_pc = pc - pc_offset;
    auto decoded = fde.decode();
    const auto& row_it = decoded.rows.find(translated_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;
}

DwarfInterpret* DwarfInterpret::get_responsible_instance(uintptr_t pc) const {
    if(map_entry.mem_region.begin <= pc && pc < map_entry.mem_region.end)
        return nullptr;
    return &acquire(pc);
}