#include "loader.h"
#include "file_reader.h"
#include "string_reader.h"
#include "machine.h"
#include "constants.h"
#include <algorithm>
#include <cstdio>

Loader::~Loader()
{
    _machine.reset();
}




void Loader::load()
{
    HeaderMetadata header = readHeader();
    _relocation_address = header.start_address;
      
    while(true) {
        RecordType type = parseRecordType(static_cast<char>(_file_reader->readByte()));
        switch (type) {
            case RecordType::TEXT: {
                TextRecord textRecord = readTextRecord();
                if (!load_into_memory(textRecord.start_address, textRecord.data)) {
                    throw std::runtime_error("Failed to load text record into memory");
                }
                break;
            }
            case RecordType::MODIFICATION: {
                ModificationRecord modRecord = readModificationRecord();
                applyModification(modRecord);
                break;
            }
            case RecordType::END: {
                EndRecord endRecord = readEndRecord();
                _machine->setPC(endRecord.execution_start_address);
                return; // Loading complete
            }
            case RecordType::UNKNOWN:
            default:
                throw std::runtime_error("Unknown record type encountered");
        }
    }
}

Loader::RecordType Loader::parseRecordType(char c)
{
    switch (c) {
        case 'H': return RecordType::HEADER;
        case 'T': return RecordType::TEXT;
        case 'M': return RecordType::MODIFICATION;
        case 'E': return RecordType::END;
        default: return RecordType::UNKNOWN; // fallback; adjust as needed
    }
}

Loader::HeaderMetadata Loader::readHeader()
{
    
    RecordType type = parseRecordType(static_cast<char>(_file_reader->readByte()));
    if (type != RecordType::HEADER) {
        throw std::runtime_error("Expected HEADER record");
    }
    if(FILE_CONTAINS_WHITE_SPACES) _file_reader->readByte();

    HeaderMetadata header;
    // Read program name (6 bytes)
    header.program_name = _file_reader->readString(6);
    if(FILE_CONTAINS_WHITE_SPACES) _file_reader->readByte();

    // Read start address (6 hex digits)
    header.start_address = std::stoi(_file_reader->readString(6), nullptr, 16);
    if(FILE_CONTAINS_WHITE_SPACES) _file_reader->readByte();

    // Read length (6 hex digits)
    header.length = std::stoi(_file_reader->readString(6), nullptr, 16);
    // consume newline
    _file_reader->readLine();
    return header;
}

Loader::TextRecord Loader::readTextRecord()
{
    if(FILE_CONTAINS_WHITE_SPACES) _file_reader->readByte();
    TextRecord record;
    // Assume 'T' has already been read
    record.start_address = std::stoi(_file_reader->readString(6), nullptr, 16);
    if(FILE_CONTAINS_WHITE_SPACES) _file_reader->readByte();

    // Read length (1 byte, 2 hex digits)
    std::string lengthStr = _file_reader->readString(2);
    int length = std::stoi(lengthStr, nullptr, 16);
    
    // Read the rest of the line (data bytes with spaces)
    std::string dataLine = _file_reader->readLine();
    
    // Remove all whitespace from the data line
    dataLine.erase(std::remove_if(dataLine.begin(), dataLine.end(), ::isspace), dataLine.end());
    
    // Now use StringReader to parse the hex bytes
    StringReader stringReader(dataLine);
    record.data.resize(length);
    
    for (int i = 0; i < length; ++i) {
        std::string byteHex = stringReader.readString(2);
        record.data[i] = static_cast<uint8_t>(std::stoi(byteHex, nullptr, 16));
    }
    
    return record;
}

Loader::ModificationRecord Loader::readModificationRecord()
{
    ModificationRecord record;
    if(FILE_CONTAINS_WHITE_SPACES) _file_reader->readByte();
    
    record.address = std::stoi(_file_reader->readString(6), nullptr, 16);
    if(FILE_CONTAINS_WHITE_SPACES) _file_reader->readByte();
    
    record.length = std::stoi(_file_reader->readString(2), nullptr, 16);
    
    record.add = true;
    std::string rest = _file_reader->readLine();
    // Remove whitespace
    rest.erase(std::remove_if(rest.begin(), rest.end(), ::isspace), rest.end());
    if (!rest.empty()) {
        if (rest[0] == '-') {
            record.add = false;
        }
    }
    
    return record;
}

Loader::EndRecord Loader::readEndRecord()
{
    EndRecord record;
    if(FILE_CONTAINS_WHITE_SPACES) _file_reader->readByte();
    // Assume 'E' has already been read
    std::string addrStr = _file_reader->readString(6);
    if (!addrStr.empty()) {
        record.execution_start_address = std::stoi(addrStr, nullptr, 16);
    } else {
        record.execution_start_address = 0; 
    }
    // consume newline
    _file_reader->readLine();
    return record;
}

bool Loader::load_into_memory(int start_address, const std::vector<uint8_t> &data)
{
    for(size_t i = 0; i < data.size(); ++i) {
        int addr = start_address + static_cast<int>(i);
        if (addr < 0 || addr >= MEMORY_SIZE) {
            return false; 
        }
        _machine->setByte(addr, data[i]);
    }
    return true;
}

void Loader::applyModification(const ModificationRecord& mod)
{
    // M record specifies address and length in half-bytes (nibbles)
    // We need to modify the value at that address by adding or subtracting
    // the relocation address
    
    int address = mod.address;
    int halfBytes = mod.length;
    
    // Calculate how many full bytes we need to read
    // halfBytes can be odd or even
    int numBytes = (halfBytes + 1) / 2;
    
    if (address < 0 || address + numBytes > MEMORY_SIZE) {
        throw std::runtime_error("Modification address out of bounds");
    }
    
    // Read the current value from memory
    int currentValue = 0;
    for (int i = 0; i < numBytes; ++i) {
        currentValue = (currentValue << 8) | _machine->getByte(address + i);
    }
    
    // If odd number of half-bytes, we only modify the relevant nibbles
    // For simplicity, we'll work with the full bytes and mask appropriately
    int mask = 0;
    for (int i = 0; i < halfBytes; ++i) {
        mask = (mask << 4) | 0xF;
    }
    
    // Extract the value to modify
    int shift = (numBytes * 2 - halfBytes) * 4;
    int valueToModify = (currentValue >> shift) & mask;
    
    // Apply modification
    int newValue = mod.add ? (valueToModify + _relocation_address) 
                           : (valueToModify - _relocation_address);
    
    // Mask to keep only the relevant bits
    newValue &= mask;
    
    // Reconstruct the full value
    int preservedBits = currentValue & ~(mask << shift);
    int finalValue = preservedBits | (newValue << shift);
    
    // Write back to memory (big-endian)
    for (int i = 0; i < numBytes; ++i) {
        _machine->setByte(address + i, (finalValue >> ((numBytes - 1 - i) * 8)) & 0xFF);
    }
}