#include "segment.h"

#include "operators.h"

#include <boost/endian/conversion.hpp>

#include <iostream>
#include <exception>

using namespace std::string_literals;

size_t Segment::getAddressOfLabel(const std::string& label)
{
 size_t address{0};
 auto i{begin()};
 while (i != end()) {
  Chunk& chunk{**i};
  address += chunk.size();
  if (typeid(chunk) == typeid(Label)) {
   if (dynamic_cast<Label&>(chunk).name() == label) {
    return address;
   }
  }
  i++;
 }

 throw std::runtime_error("Error: Label |"s + label + "| not found."s);
}

size_t Segment::getAddressOfIndex(size_t index)
{
 if (index >= size())
  throw std::runtime_error("Error: Index |"s + std::to_string(index) + "| not found."s);

 size_t address{0};
 for (size_t i = 0; i < index; i++) {
  address += (*this)[i]->size();
 }
 return address;
}

std::vector<uint8_t> Segment::getCode()
{
 std::vector<uint8_t> result;

 for (const auto& chunk: *this)
  result = result + chunk->getCode();

 return result;
}

void Segment::insertAddresses()
{
 // insert relative addresses
 for (size_t i = 0; i < this->size(); i++) {
  try {
   AddressFeature& af { dynamic_cast<AddressFeature&>(*((*this)[i]))};

   if (af.relativeAddressing == false)
    throw std::runtime_error("Error: Absolute Addressing not supported.");

   int64_t target_address = getAddressOfLabel(af.label); // throws if label not found
   int64_t start_address = getAddressOfIndex(i); // throws if out of range
   int64_t diff = target_address - start_address;
   uint64_t udiff = static_cast<uint64_t>(diff);
   if (af.addr_size == 1) {
    if (diff < -128 || diff > 127)
     throw std::runtime_error("Error: Address too far."s);
    af.machine_code[af.addr_offs] = static_cast<uint8_t>(udiff);
   } else if (af.addr_size == 4) {
    *(reinterpret_cast<uint32_t*>(af.machine_code.data() + af.addr_offs)) = boost::endian::native_to_little(static_cast<uint32_t>(udiff));
   } else
    throw std::runtime_error("Error: unexpected addr_size: "s + std::to_string(af.addr_size));
  } catch (const std::bad_cast& ex) {
   // ignore, expected for non-addressing chunks
  }
 }
}

void Segment::optimize()
{
 bool changed{false};

 for (size_t i = 0; i < this->size(); i++) {
  try {
   AddressFeature& af { dynamic_cast<AddressFeature&>(*((*this)[i]))};

   if (af.relativeAddressing == false)
    throw std::runtime_error("Error: Absolute Addressing not supported.");

   if (!af.alternative_code.empty()) {
    int64_t target_address = getAddressOfLabel(af.label); // throws if label not found
    int64_t start_address = getAddressOfIndex(i); // throws if out of range
    int64_t diff = target_address - start_address;
    if (af.addr_size == 4 && af.alternative_size == 1 && diff >= -128 && diff <= 127) {
     af.machine_code = af.alternative_code;
     af.addr_size = af.alternative_size;
     af.addr_offs = af.alternative_offs;
     af.alternative_code.clear();
     af.alternative_size = 0;
     af.alternative_offs = 0;
     changed = true;
    }
   }
  } catch (const std::bad_cast& ex) {
   // ignore, expected for non-addressing chunks
  }
 }

 if (changed)
  insertAddresses(); // renumber
}