/* 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;
        case 'r':  return 0x1B;
        case ':':  return 0x1F; // replaces !
        case 'i':  return 0x1E; // replaces '
        case (unsigned char)'\xC4': return 0x1D; // Auml replaces ,
        //case (unsigned char)'\xD6': return 0x21; // Ouml replaces '
        default: return c - 0x40;
    }
}
static unsigned char TranslateOtherChar(unsigned char c)
{
    switch(c)
    {
        case '0' ... '9': return (c-'0') + 0xA0;
        case ' ': return 0x40;
        case '-': return 0x94;
        case '!': return 0x5F;
        case '.': return 0x5C;
        case '\'': return 0x5E;
        case ',': return 0x5D;
        case (unsigned char)'\xC4': return 0x5B; // Auml
        case (unsigned char)'\xD6': return 0x5E; // Ouml
        default: return c;
    }
}
static unsigned char TranslateIntroChar(unsigned char c)
{
    switch(c)
    {
        case 'c': return 0xA3;
        case '0': return 0xA0;
        case '1': return 0xA7;
        case '2': return 0xA2;
        case '9': return 0xA5;
        case '8': return 0xA1; case '4': return 0x9F;
        case '.': return 0xDC;
        case ',': return 0xDD;
        case ' ': return 0x00;
        case '-': return 0x80;
        case ':': return 0x83;
        case 'j': return 0xA8; case '+': return 0xA9;
        case '/': return 0xAA;
/*
        // Special lower-case letters that are only used in the URL
        case 'b': return 0x8D;
        case 'h': return 0x8A;
        case 'i': return 0x8B;
        case 'k': return 0x84;
        case 'p': return 0x85;
        case 'q': return 0x86;
        case 's': return 0x87;
        case 't': return 0x88;
        case 'w': return 0x89;
        case 'f': return 0x8C;
        case 'm': return 0x8E;
        case 'n': return 0x8F;*/
        case (unsigned char)'\xC4': return 0x81; // Auml
        case (unsigned char)'\xD6': return 0x82; // Ouml
        default: return c+0x80;
    }
}

template<typename Translator>
static std::string TranslateString(const std::string& str, Translator&& Translate)
{
    std::string result;
    std::regex pat("\\[(PPU )?([0-9A-F]+)\\].*");
    for(unsigned a=0; a<str.size(); ++a)
    {
        std::smatch res;
        if(std::regex_match(str.cbegin()+a, str.cend(), res, pat))
        {
            //fprintf(stderr, "res1(%s)res2(%s)\n", std::string(res[1]).c_str(), std::string(res[2]).c_str());
            if(res[1] == "PPU ")
            {
                unsigned ppuaddr = std::stoi(res[2], nullptr, 16);
                result += char(ppuaddr>>8);
                result += char(ppuaddr&0xFF);
            }
            else
                result += char( std::stoi(res[2], nullptr, 16) );
            a = res[2].second - str.begin();
        }
        else if(str[a] == '\\')
            result += Translate(str[++a]);
        else
            result += Translate(str[a]);
    }
    return result;
}

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 RefType { lowbyte, highbyte, word, length, pickbyte } 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 AddRef(const std::string& targetname, unsigned where, Reference::RefType reftype, unsigned offset=0)
    {
        refs.emplace_back(Reference{targetname,offset, where, reftype});
    }
    void AddFreeSpace(unsigned address, unsigned length)
    {
        space[address & 0x3C000].emplace_back(address,length);
    }
    void Link()
    {
        std::map<std::string, std::pair<unsigned/*where*/,unsigned/*size*/>> addresses;
        // Sort the blobs in order of length (largest first)
        // This helps the space-fitting algorithm a bit
        std::sort(blobs.begin(), blobs.end(), [](const Blob& a, const Blob& b)
            { return a.data.size() > b.data.size(); });
        for(auto& b: blobs)
        {
            if(b.fixed_addr == false && b.bank != ~0u)
            {
                unsigned chose=~0u, chosesize=~0u;
                for(auto& s2: space[b.bank])
                    if(s2.second >= b.data.size())
                        if(s2.second < chosesize)
                            { chosesize = s2.second; chose = s2.first; }
                if(chose != ~0u)
                    for(auto& s2: space[b.bank])
                        if(s2.first == chose)
                        {
                            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;
            }
            addresses[b.name] = std::make_pair(b.fixed_addr ? b.address : ~0u, b.data.size());
        }
        for(const auto& r: refs)
        {
            auto i = addresses.find(r.name);
            if(i == addresses.end())
            {
                std::cerr << "ERROR1: Could not generate reference to " << r.name << " -- symbol is nowhere\n";
                continue;
            }
            if(i->second.first == ~0u && r.type != Reference::length)
            {
                std::cerr << "ERROR2: Could not generate reference to " << r.name << " -- symbol is nowhere\n";
                continue;
            }
            Blob b;
            b.fixed_addr = true;
            b.address    = r.address;
            b.name = GenerateName("REF_TO_" + r.name);
            unsigned target = ROM2NESaddr(i->second.first + 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;
                case Reference::length:
                    b.data += char(i->second.second + r.offset);
                    break;
                case Reference::pickbyte:
                {
                    for(const auto& b1: blobs)
                        if(b1.name == r.name)
                        {
                            b.data += b1.data[r.offset];
                            break;
                        }
                    break;
                }
                default:
                    std::cerr << "Unknown reference type: " << r.type << std::endl;
            }
            blobs.push_back(std::move(b));
        }
    }
    void Dump()
    {
        // Sort the blobs in name order
        std::sort(blobs.begin(), blobs.end(), [](const Blob& a, const Blob& b)
            { return a.name < b.name; });
        auto DumpData = [](const std::string& s)
        {
            std::string result;
            for(unsigned char c: s)
            {
                char Buf[8];
                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","length","pickbyte"};
            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";
        // Sort the blobs in order of address
        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,unsigned>>&& tokens,
                    std::map<unsigned, std::vector<char>>&& font)
{
    Linker linker;
    Linker::Blob currenttable, currentstring;
    enum { nothing, screentable, introtable, fstring, ending } mode = nothing;
    unsigned fixedwidth = 0;
    auto FlushTable = [&]
    {
        switch(mode)
        {
            case nothing: break;
            case fstring: break;
            case screentable:
                linker.AddBlob(std::move(currentstring));
                linker.AddBlob(std::move(currenttable));
                break;
            case introtable:
                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();
    };
    /* Token numbers come from the regex in main.
     * 0 = comment
     * 1 = alphanumeric, optionally prefixed with minus
     * 2 = string delimited with ""
     * 3 = string delimited with ''
     * 4 = string delimited with <>
     */
    for(unsigned a=0; a<tokens.size(); ++a)
    {
        const auto& t = tokens[a];
        if(t.second==1)
        {
            if(t.first == "SCREENTABLE")
            {
                FlushTable();
                mode               = screentable;
                currentstring.name = linker.GenerateName("SCREENDATA");
                currenttable.name  = linker.GenerateName("SCREENOFFSET");
                currentstring.fixed_addr = false;
                currenttable.fixed_addr  = false;
                currentstring.bank       = 0x34000;
                currenttable.bank        = 0x34000;
            }
            if(t.first == "INTROTABLE")
            {
                FlushTable();
                mode               = introtable;
                currentstring.name = linker.GenerateName("INTRODATA");
                currentstring.fixed_addr = false;
                currentstring.bank       = 0x34000;
                currenttable.name = linker.GenerateName("INTROLENGTH_DUMMY");
                currenttable.fixed_addr = false;
                currenttable.bank       = ~0u;
            }
            if(t.first == "REFERRED_FROM")
            {
                unsigned target = stoi(tokens[a+1].first, nullptr, 16);
                unsigned type   = stoi(tokens[a+2].first);
                a += 2;
                switch(type & 0xF)
                {
                    // References to string table
                    case 0: linker.AddRef(currentstring.name, target, Linker::Reference::lowbyte,  char(type>>4)); break;
                    case 1: linker.AddRef(currentstring.name, target, Linker::Reference::highbyte, char(type>>4)); break;
                    case 2: linker.AddRef(currentstring.name, target, Linker::Reference::word,     char(type>>4)); break;
                    case 3: linker.AddRef(currentstring.name, target, Linker::Reference::length,   char(type>>4)); break;
                    // References to pointer table
                    case 4: linker.AddRef(currenttable.name,  target, Linker::Reference::lowbyte,  char(type>>4)); break;
                    case 5: linker.AddRef(currenttable.name,  target, Linker::Reference::highbyte, char(type>>4)); break;
                    case 6: linker.AddRef(currenttable.name,  target, Linker::Reference::word,     char(type>>4)); break;
                    case 7: linker.AddRef(currenttable.name,  target, Linker::Reference::length,   char(type>>4)); break;
                    // _Contents_ from pointer table
                    case 8: linker.AddRef(currenttable.name,  target, Linker::Reference::pickbyte, type>>4); break;

                    default: std::cerr << "Unknown refer type " << type << std::endl;
                }
            }
            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")
            {
                unsigned address = stoi(tokens[a+1].first, nullptr, 16);
                a += 1;
                FlushTable();
                mode = ending;
                currentstring.name = linker.GenerateName("ENDINGDATA");
                currenttable.name = linker.GenerateName("ENDINGLENGTHTABLE");
                if(address)
                {
                    currentstring.fixed_addr = false;
                    currentstring.bank       = address & 0x3C000;
                    currenttable.fixed_addr  = true;
                    currenttable.address     = address;
                }
                else
                {
                    currentstring.fixed_addr = false;
                    currentstring.bank       = 0x24000;
                    currenttable.fixed_addr  = false;
                    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
        {
            // Actions 
            std::string word;
            switch(t.second)
            {
                case 2: word = TranslateString(t.first, TranslateOtherChar); break;
                case 3: word = TranslateString(t.first, TranslateEndingChar); break;
                case 4: word = TranslateString(t.first, TranslateIntroChar); break;
            }
            switch(mode)
            {
                case screentable:
                {
                    word.insert(word.begin()+2, char(word.size()-2));
                    auto pos = currentstring.data.find(word);
                    if(pos != currentstring.data.npos)
                    {
                        std::fprintf(stderr, "Screentable: String %s reused\n", t.first.c_str());
                        currenttable.data.push_back(pos);
                        return;
                    }
                    else
                    {
                        currenttable.data.push_back(currentstring.data.size());
                    }
                    break;
                }
                case introtable:
                    word.insert(word.begin(), char(word.size()-2));
                    currenttable.data.push_back('\0');
                    break;
                case ending:
                    currenttable.data.push_back(word.size());
                    break;
                default: break;
            }
            currentstring.data += word;
            switch(mode)
            {
                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;
                }
                default: break;
            }
        }
    }
    for(const auto& f: font)
    {
        Linker::Blob ch;
        ch.name       = "font" + std::to_string(f.first);
        ch.fixed_addr = true;
        ch.address    = f.first;
        ch.data = std::string(16,'\0');
        if(f.second.size() != 64)
            std::cerr << "Invalid syntax in font character " << std::hex << f.first << std::dec << std::endl;
        else
        {
            for(unsigned y=0; y<8; ++y)
                for(unsigned x=0; x<8; ++x)
                {
                    char c = f.second[y*8+x];
                    unsigned v = 0;
                    if(c=='#') v = 1; else if(c == '.') v = 2; else if(c == 'c') v = 3;
                    ch.data[y+0] |= ((v &1) << (7-x));
                    ch.data[y+8] |= ((v>>1) << (7-x));
                }
        }
        linker.AddBlob(std::move(ch));
    }
    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,unsigned/*type*/>> tokens;
    std::map<unsigned/*address*/,std::vector<char>/*repr*/> font;
    std::vector<unsigned> font_header;
    while(f.good())
    {
        std::string s;
        std::getline(f, s);
        std::smatch res;
        if(s[0] == '$')
        {
            font_header.clear();
            for(unsigned c=0; c+9<=s.size(); c+=9)
                if(s[c] == '$')
                    font_header.push_back(std::stoi(s.substr(c+1,8),nullptr,16));
        }
        else if(s[0] == ':')
        {
            for(unsigned c=0; c+9<=s.size(); c+=9)
                if(s[c] == ':')
                {
                    auto& t = font[ font_header[c/9] ];
                    t.insert(t.end(), &s[c+1], &s[c+9]);
                }
        }
        else
        {
            for(auto b = s.cbegin(), e = s.cend(); std::regex_search(b,e, res, pat); b = res[0].second)
                if(!res[1].length()) // comment
                {
                    unsigned index = 2;
                    while(res[index].length() == 0) ++index;
                    tokens.emplace_back(res[index], index-1);
                }
            tokens.emplace_back("",true);
        }
    }
    /*for(auto t: tokens)
        std::fprintf(stderr, "<%s> %u\n", t.first.c_str(), t.second);*/
    Compile(std::move(tokens), std::move(font));
    return 0;
}