/* Copyright (C) 2016 Joel Yliluoma - http://iki.fi/bisqwit/
 * Redistribution and modifications are permitted under the rules
 * of Creative Commons Attribution-ShareAlike 3.0 license
 * https://creativecommons.org/licenses/by-sa/3.0/
 */
#include <fstream>
#include <string>
#include <regex>
#include <iostream>
#include <cstdio>
#include <map>

static unsigned char TranslateEndingChar(unsigned char c)
{
    switch(c)
    {
        case '0' ... '9': return c;
        case '.':  return 0x1C;
        case '\'': return 0x1E;
        case ' ':  return 0x20;
        default: return c - 0x40;
    }
}
static unsigned char TranslateOtherChar(unsigned char c)
{
    switch(c)
    {
        case '0' ... '9': return c + 0xA0;
        case ' ': return 0x40;
        case '-': return 0x94;
        case '!': return 0x5F;
        case '.': return 0x5C;
        default: return c;
    }
}
static unsigned ROM2NESaddr(unsigned romaddr)
{
    if(romaddr >= 0x3C000) return (romaddr & 0x3FFF) | 0xC000;
    return (romaddr & 0x3FFF) | 0x8000;
}

class Linker
{
public:
    struct Blob
    {
        unsigned bank=0;    // If fixed_addr=false
        unsigned address=0; // If fixed_addr=true
        bool     fixed_addr=false;
        std::string name{};
        std::string data{};
    };
    struct Reference
    {
        std::string name;    // Target symbol that is being referred
        unsigned    offset;  // Offset to the target that is being referred
        unsigned    address; // Where is being referred from
        enum { lowbyte, highbyte, word } type;
    };
private:
    std::vector<Blob>      blobs;
    std::vector<Reference> refs;
    std::map<unsigned/*bank*/, std::vector<std::pair<unsigned/*offset*/,unsigned/*length*/>>> space;
    unsigned namecounter=0;
public:
    std::string GenerateName(const std::string& templ = "STR")
    {
        return templ + "_" + std::to_string(namecounter++);
    }
    void AddBlob(Blob&& blob)
    {
        blobs.emplace_back(std::move(blob));
    }
    void AddLowRef(const std::string& targetname, unsigned where, unsigned offset=0)
    {
        refs.emplace_back(Reference{targetname,offset, where,Reference::lowbyte});
    }
    void AddHighRef(const std::string& targetname, unsigned where, unsigned offset=0)
    {
        refs.emplace_back(Reference{targetname,offset, where,Reference::highbyte});
    }
    void AddWordRef(const std::string& targetname, unsigned where, unsigned offset=0)
    {
        refs.emplace_back(Reference{targetname,offset, where,Reference::word});
    }
    void AddFreeSpace(unsigned address, unsigned length)
    {
        space[address & 0x3C000].emplace_back(address,length);
    }
    void Link()
    {
        std::map<std::string, unsigned> addresses;
        for(auto& b: blobs)
        {
            if(b.fixed_addr == false)
            {
                for(auto& s2: space[b.bank])
                    if(s2.second >= b.data.size())
                    {
                        b.address = s2.first;
                        b.fixed_addr = true;
                        s2.first  += b.data.size();
                        s2.second -= b.data.size();
                        break;
                    }
                if(!b.fixed_addr)
                    std::cerr << "ERROR: Could not " << b.data.size()
                              << " bytes of free space in bank " << std::hex << b.bank << std::dec
                              << " for " << b.name << std::endl;
                std::cerr << "Bank contains:";
                for(auto& s2: space[b.bank])
                    std::cerr << std::hex << ' ' << s2.first << '+' << std::dec << s2.second;
                std::cerr << std::endl;
            }
            if(b.fixed_addr)
                addresses[b.name] = b.address;
        }
        for(const auto& r: refs)
        {
            auto i = addresses.find(r.name);
            if(i == addresses.end())
            {
                std::cerr << "ERROR: Could not generate reference to " << r.name << " -- symbol is nowhere\n";
                continue;
            }
            Blob b;
            b.fixed_addr = true;
            b.address    = r.address;
            unsigned target = ROM2NESaddr(i->second + r.offset);
            switch(r.type)
            {
                case Reference::lowbyte:
                    b.data += char(target & 0xFF);
                    break;
                case Reference::highbyte:
                    b.data += char(target >> 8);
                    break;
                case Reference::word:
                    b.data += char(target & 0xFF);
                    b.data += char(target >> 8);
                    break;
            }
            blobs.push_back(std::move(b));
        }
    }
    void Dump()
    {
        auto DumpData = [](const std::string& s)
        {
            std::string result;
            for(unsigned char c: s)
            {
                char Buf[4];
                std::sprintf(Buf, " %02X", c);
                result += Buf;
            }
            return result;
        };
        std::cerr << "Linker contains:\n";
        for(const auto& b: blobs)
        {
            std::cerr << "\"" << b.name << "\" in ";
            if(b.fixed_addr)
                std::cerr << " address " << std::hex << b.address << std::dec;
            else
                std::cerr << " bank    " << std::hex << b.bank    << std::dec;
            std::cerr << " data " << DumpData(b.data) << "\n";
        }
        for(const auto& r: refs)
        {
            static const char*const types[] = {"lowbyte","highbyte","word"};
            std::cerr << "Reference to \"" << r.name << "\" + " << r.offset
                      << std::hex << " from " << r.address << std::dec
                      << ", type: " << types[r.type] << std::endl;
        }
    }
    void MakePatch()
    {
        std::cout << "PATCH";
        std::sort(blobs.begin(), blobs.end(), [](const Blob& a, const Blob& b)
        {
            return a.address < b.address;
        });
        for(const auto& b: blobs)
            if(b.fixed_addr)
            {
                std::cout << char( ((b.address+16) >> 16) & 0xFF );
                std::cout << char( ((b.address+16) >>  8) & 0xFF );
                std::cout << char( ((b.address+16) >>  0) & 0xFF );
                std::cout << char( (b.data.size() >>  8) & 0xFF );
                std::cout << char( (b.data.size() >>  0) & 0xFF );
                std::cout << b.data;
            }
        std::cout << "EOF";
    }
};

static void Compile(std::vector<std::pair<std::string,bool>>&& tokens)
{
    Linker linker;
    Linker::Blob currenttable, currentstring;
    enum { nothing, screentable, fstring, ending } mode = nothing;
    unsigned fixedwidth = 0;
    auto FlushTable = [&]
    {
        switch(mode)
        {
            case nothing: break;
            case fstring: break;
            case screentable:
                currenttable.data.erase(currenttable.data.size()-1);
                linker.AddBlob(std::move(currentstring));
                linker.AddBlob(std::move(currenttable));
                break;
            case ending:
                linker.AddBlob(std::move(currentstring));
                linker.AddBlob(std::move(currenttable));
                break;
        }
        currenttable  = Linker::Blob();
    };
    auto FlushElement = [&]
    {
        switch(mode)
        {
            case nothing: break;
            case fstring:
            {
                if(currentstring.data.size() > fixedwidth)
                {
                    std::cerr << "Fixed string \"" << currentstring.data << "\" too long!\n";
                }
                else
                    currentstring.data.append(fixedwidth - currentstring.data.size(), 0x40);
                linker.AddBlob(std::move(currentstring));
                break;
            }
            case ending:
            {
                break;
            }
            case screentable:
            {
                currenttable.data += currentstring.data.size(); // Offset of the next string
                break;
            }
        }
    };
    for(unsigned a=0; a<tokens.size(); ++a)
    {
        const auto& t = tokens[a];
        if(t.second)
        {
            if(t.first == "PPU")
            {
                unsigned target = stoi(tokens[a+1].first, nullptr, 16);
                currentstring.data += char(target >> 8);
                currentstring.data += char(target & 0xFF);
                continue;
            }
            if(t.first == "SCREENTABLE")
            {
                FlushTable();
                mode               = screentable;
                currentstring.name = linker.GenerateName("SCREENDATA");
                currenttable.name  = linker.GenerateName("SCREENOFFSET");
                currenttable.data += '\0'; // Offset of the first string
                currentstring.fixed_addr = false;
                currenttable.fixed_addr  = false;
                currentstring.bank       = 0x34000;
                currenttable.bank        = 0x34000;
            }
            if(t.first == "REFER")
            {
                unsigned target = stoi(tokens[a+1].first, nullptr, 16);
                unsigned type   = stoi(tokens[a+2].first);
                a += 2;
                switch(mode)
                {
                    case nothing: // N/A
                    case fstring: // N/A
                        break;
                    case screentable:
                        if(type == 2)
                            linker.AddWordRef(currenttable.name, target);
                        else
                            linker.AddWordRef(currentstring.name, target, type);
                        break;
                    case ending:
                        switch(type)
                        {
                            case 0: linker.AddHighRef(currentstring.name, target); break;
                            case 1: linker.AddLowRef(currentstring.name, target); break;
                            case 2: linker.AddWordRef(currenttable.name, target); break;
                            default:
                                std::cerr << "Unknown patch type " << type << "\n";
                        }
                        break;
                }
            }
            if(t.first == "FSTRING")
            {
                FlushTable();
                unsigned length  = stoi(tokens[a+1].first);
                unsigned address = stoi(tokens[a+2].first, nullptr, 16);
                a += 2;
                mode = fstring;
                currentstring.name       = linker.GenerateName("FSTRING");
                currentstring.address    = address;
                currentstring.fixed_addr = true;
                fixedwidth = length;
            }
            if(t.first == "ENDINGTABLE")
            {
                FlushTable();
                mode = ending;
                currentstring.name = linker.GenerateName("ENDINGDATA");
                currenttable.name = linker.GenerateName("ENDINGOFFSET");
                currentstring.fixed_addr = false;
                currenttable.fixed_addr  = false;
                currentstring.bank       = 0x24000;
                currenttable.bank        = 0x24000;
            }
            if(t.first == "FREE")
            {
                unsigned address = stoi(tokens[a+1].first, nullptr, 16);
                unsigned length  = stoi(tokens[a+2].first);
                a += 2;
                linker.AddFreeSpace(address, length);
            }
        }
        else
        {
            if(mode == screentable)
                currentstring.data.push_back(t.first.size());
            if(mode == ending)
                currenttable.data.push_back(t.first.size());
            for(unsigned char c: t.first)
            {
                c = (mode == ending)
                    ? TranslateEndingChar(c)
                    : TranslateOtherChar(c);
                currentstring.data.push_back(c);
            }
            /*if(mode == screentable)*/ FlushElement();
        }
    }
    FlushTable();
    linker.Link();
    linker.MakePatch();
    linker.Dump();
}

int main(int argc, char**argv)
{
    std::regex pat("([A-Z0-9]+)|\"([^\"]*)\"|(#.*)");
    std::ifstream f(argv[1]);
    std::vector<std::pair<std::string,bool>> tokens;
    while(f.good())
    {
        std::string s;
        std::getline(f, s);
        std::smatch res;
        for(auto b = s.cbegin(), e = s.cend(); std::regex_search(b,e, res, pat); b = res[0].second)
            if(!res[3].length())
                tokens.emplace_back(res[1].length() ? res[1] : res[2], !!res[1].length());
        tokens.emplace_back("",true);
    }
    Compile(std::move(tokens));
    return 0;
}