module jpeg;

private
{
    import std.stdio;
    import std.file;
    import std.algorithm;
    import std.conv;
    import dlib.math.utils;
    import bitio;
    import huffman;
}

T readNumeric(T) (File* f, Endian endian = Endian.Little) if (is(T == ubyte))
{
    ubyte[1] bytes;
    f.rawRead(bytes);
    return bytes[0];
}

T readNumeric(T) (File* f, Endian endian = Endian.Little) if (is(T == ushort))
{
    union U16
    {
        ubyte[2] asBytes;
        ushort asUshort;
    };
    U16 u16;
    f.rawRead(u16.asBytes);
    version(LittleEndian)
    {
        if (endian == Endian.Big)
            return u16.asUshort.swapEndian16;
        else
            return u16.asUshort;
    }
    else
    {
        if (endian == Endian.Little)
            return u16.asUshort.swapEndian16;
        else
            return u16.asUshort;
    }
}

char[Size] readChars(size_t Size) (File* f)
{
    char[Size] chars;
    f.rawRead(chars);
    return chars;
}

/*
    TODO:
    1. Add JPEGImage class to store everything
    (width, height, color depth, Huffman tables,
    quantization matrices, image data etc.)
    Pass this image when reading markers.
    Get rid of the redundant JPEGMarker* objects -
    readMarker function should return only
    marker type and update JPEGImage.
    This will greatly reduce amount of code.
    2. Decode SOS compressed data using Huffman
    tables. Learn how image channels (components)
    are separated from each other.
*/

class JPEGImage
{
    struct JFIF
    {
        ubyte versionMajor;
        ubyte versionMinor;
        ubyte units;
        ushort xdensity;
        ushort ydensity;
        ubyte thumbnailWidth;
        ubyte thumbnailHeight;
        ubyte[] thumbnail;
    }

    struct DQT
    {
        ubyte precision;
        ubyte tableId;
        ubyte[] table;
    }

    struct SOF0Component
    {
        ubyte hsubsamling;
        ubyte vsubsamling;
        ubyte tableId;
    }

    struct SOF0
    {
        ubyte precision;
        ushort height;
        ushort width;
        ubyte componentsNum;
        SOF0Component[] components;
    }

    struct DHT
    {
        ubyte clas;
        ubyte tableId;
        ubyte[HuffmanCode] huffmanTable;
        HuffmanTreeNode* huffmanTree;
    }

    struct SOSComponent
    {
        ubyte tableIdDC;
        ubyte tableIdAC;
    }

    struct SOS
    {
        ubyte componentsNum;
        SOSComponent[] components;
        ubyte spectralSelectionStart;
        ubyte spectralSelectionEnd;
        ubyte successiveApproximationBitHigh;
        ubyte successiveApproximationBitLow;
        ubyte[] encodedData;
    }

    JFIF jfif;
    DQT[] dqt;
    SOF0 sof0;
    DHT[] dht;
    SOS sos;

    DQT* getQuantizationTable(ubyte id)
    {
        foreach(ref t; dqt)
            if (t.tableId == id)
                return &t;
        return null;
    }

    DHT* getHuffmanTable(ubyte clas, ubyte id)
    {
        foreach(ref t; dht)
            if (t.clas == clas &&
                t.tableId == id)
                return &t;
        return null;
    }
}

enum JPEGMarkerType
{
    Unknown,
    SOI,
    SOF0,
    SOF1,
    SOF2,
    DHT,
    DQT,
    DRI,
    SOS,
    RSTn,
    APP0,
    COM,
    EOI
}

void loadJPEG(string filename)
{
    auto jpg = new JPEGImage();
    auto f = new File(filename, "r");

    while (!f.eof)
    {
        JPEGMarkerType mt = f.readMarker(jpg);
    }

    f.close();

    decodeScanData(jpg);
}

JPEGMarkerType readMarker(File* f, JPEGImage jpg)
{
    JPEGMarkerType mt;
    ushort magic = f.readNumeric!ushort(Endian.Big);

    switch (magic)
    {
        case 0xFFD8:
            mt = JPEGMarkerType.SOI;
            break;

        case 0xFFE0:
            mt = JPEGMarkerType.APP0;
            readJFIF(f, jpg);
            break;

        case 0xFFDB:
            mt = JPEGMarkerType.DQT;
            readDQT(f, jpg);
            break;

        case 0xFFC0:
            mt = JPEGMarkerType.SOF0;
            readSOF0(f, jpg);
            break;

        case 0xFFC4:
            mt = JPEGMarkerType.DHT;
            readDHT(f, jpg);
            break;

        case 0xFFDA:
            mt = JPEGMarkerType.SOS;
            readSOS(f, jpg);
            break;

        default:
            break;
    }

    return mt;
}

void readJFIF(File* f, JPEGImage jpg)
{
    ushort jfif_length = f.readNumeric!ushort(Endian.Big);

    char[5] jfif_id = f.readChars!5;
    assert(jfif_id == "JFIF\0");

    jpg.jfif.versionMajor = f.readNumeric!ubyte;
    jpg.jfif.versionMinor = f.readNumeric!ubyte;
    jpg.jfif.units = f.readNumeric!ubyte;
    jpg.jfif.xdensity = f.readNumeric!ushort(Endian.Big);
    jpg.jfif.ydensity = f.readNumeric!ushort(Endian.Big);
    jpg.jfif.thumbnailWidth = f.readNumeric!ubyte;
    jpg.jfif.thumbnailHeight = f.readNumeric!ubyte;

    uint jfif_thumb_length = jpg.jfif.thumbnailWidth * jpg.jfif.thumbnailHeight * 3;
    if (jfif_thumb_length > 0)
    {
        jpg.jfif.thumbnail = new ubyte[jfif_thumb_length];
        f.rawRead(jpg.jfif.thumbnail);
    }

    writefln("APP0/JFIF length: %s", jfif_length);
    writefln("APP0/JFIF identifier: %s", jfif_id);
    writefln("APP0/JFIF version major: %s", jpg.jfif.versionMajor);
    writefln("APP0/JFIF version minor: %s", jpg.jfif.versionMinor);
    writefln("APP0/JFIF units: %s", jpg.jfif.units);
    writefln("APP0/JFIF xdensity: %s", jpg.jfif.xdensity);
    writefln("APP0/JFIF ydensity: %s", jpg.jfif.ydensity);
    writefln("APP0/JFIF xthumbnail: %s", jpg.jfif.thumbnailWidth);
    writefln("APP0/JFIF ythumbnail: %s", jpg.jfif.thumbnailHeight);
}

void readDQT(File* f, JPEGImage jpg)
{
    ushort dqt_length = f.readNumeric!ushort(Endian.Big);

    JPEGImage.DQT dqt;
    jpg.dqt ~= dqt;

    ubyte bite = f.readNumeric!ubyte;
    jpg.dqt[$-1].precision = bite.hiNibble;
    jpg.dqt[$-1].tableId = bite.loNibble;

    if (jpg.dqt[$-1].precision == 0)
        jpg.dqt[$-1].table = new ubyte[64];
    else if (jpg.dqt[$-1].precision == 1)
        jpg.dqt[$-1].table = new ubyte[64*2];
    f.rawRead(jpg.dqt[$-1].table);

    writefln("DQT length: %s", dqt_length);
    writefln("DQT precision: %s", jpg.dqt[$-1].precision);
    writefln("DQT table id: %s", jpg.dqt[$-1].tableId);
    writefln("DQT table: %s", jpg.dqt[$-1].table);
}

void readSOF0(File* f, JPEGImage jpg)
{
    ushort sof0_length = f.readNumeric!ushort(Endian.Big);
    jpg.sof0.precision = f.readNumeric!ubyte;
    jpg.sof0.height = f.readNumeric!ushort(Endian.Big);
    jpg.sof0.width = f.readNumeric!ushort(Endian.Big);
    jpg.sof0.componentsNum = f.readNumeric!ubyte;
    assert(jpg.sof0.componentsNum == 3);

    writefln("SOF0 length: %s", sof0_length);
    writefln("SOF0 precision: %s", jpg.sof0.precision);
    writefln("SOF0 height: %s", jpg.sof0.height);
    writefln("SOF0 width: %s", jpg.sof0.width);
    writefln("SOF0 components: %s", jpg.sof0.componentsNum);

    jpg.sof0.components = new JPEGImage.SOF0Component[jpg.sof0.componentsNum];

    foreach(ref c; jpg.sof0.components)
    {
        ubyte c_id = f.readNumeric!ubyte;
        ubyte bite = f.readNumeric!ubyte;
        c.hsubsamling = bite.hiNibble;
        c.vsubsamling = bite.loNibble;
        c.tableId = f.readNumeric!ubyte;
        writefln("SOF0 component id: %s", c_id);
        writefln("SOF0 component %s hsubsamling: %s", c_id, c.hsubsamling);
        writefln("SOF0 component %s vsubsamling: %s", c_id, c.vsubsamling);
        writefln("SOF0 component %s table id: %s", c_id, c.tableId);
    }
}

void readDHT(File* f, JPEGImage jpg)
{
    ushort dht_length = f.readNumeric!ushort(Endian.Big);
    dht_length -= 2;
    writefln("DHT length: %s", dht_length);

    while(dht_length > 0)
    {
    JPEGImage.DHT dht;
    jpg.dht ~= dht;

    ubyte bite = f.readNumeric!ubyte;
    dht_length--;
    jpg.dht[$-1].clas = bite.hiNibble;
    jpg.dht[$-1].tableId = bite.loNibble;

    ubyte[16] dht_code_lengths;
    f.rawRead(dht_code_lengths);
    dht_length -= 16;

    writefln("DHT class: %s (%s)", jpg.dht[$-1].clas, jpg.dht[$-1].clas? "AC":"DC"); // AC/DC :)
    writefln("DHT tableId: %s", jpg.dht[$-1].tableId);
    writefln("DHT Huffman code lengths: %s", dht_code_lengths);

    // Read Huffman table
    int totalCodes = reduce!("a + b")(0, dht_code_lengths);
    int storedCodes = 0;
    ubyte treeLevel = 0;
    ushort bits = 0;

    while (storedCodes != totalCodes)
    {
        while (treeLevel < 15 &&
               dht_code_lengths[treeLevel] == 0)
        {
            treeLevel++;
            bits *= 2;
        }

        if (treeLevel < 16)
        {
            uint bitsNum = treeLevel + 1;
            HuffmanCode code = HuffmanCode(bits, bitsNum);
            //ubyte symbol = f.readNumeric!ubyte;
            jpg.dht[$-1].huffmanTable[code] = f.readNumeric!ubyte;
            dht_length--;

            storedCodes++;
            bits++;
            dht_code_lengths[treeLevel]--;
        }
    }

    writeln("DHT Huffman table:\n", jpg.dht[$-1].huffmanTable);

    //writeln("DHT Huffman table:");
    //foreach(i, v; jpg.dht[$-1].huffmanTable)
    //    writefln("%s: %x", i, v);
    auto huffmanSymbols = jpg.dht[$-1].huffmanTable.values;
    huffmanSymbols.sort();
    foreach(v; huffmanSymbols)
        writef("%x ", v);
    writef("\n");

    //auto tree = jpg.dht[$-1].huffmanTree;

    jpg.dht[$-1].huffmanTree = treeFromTable(jpg.dht[$-1].huffmanTable);
    string[ubyte] table;
    jpg.dht[$-1].huffmanTree.getCodes(table);
    writeln(table);
    }
}

void readSOS(File* f, JPEGImage jpg)
{
    ushort sos_length = f.readNumeric!ushort(Endian.Big);
    jpg.sos.componentsNum = f.readNumeric!ubyte;

    writefln("SOS length: %s", sos_length);
    writefln("SOS components: %s", jpg.sos.componentsNum);

    //assert(jpg.sos.componentsNum == 3);

    jpg.sos.components = new JPEGImage.SOSComponent[jpg.sos.componentsNum];

    foreach(ref c; jpg.sos.components)
    {
        ubyte c_id = f.readNumeric!ubyte;
        ubyte bite = f.readNumeric!ubyte;
        c.tableIdDC = bite.hiNibble;
        c.tableIdAC = bite.loNibble;
        writefln("SOS component id: %s", c_id);
        writefln("SOS component %s DC table id: %s", c_id, c.tableIdDC);
        writefln("SOS component %s AC table id: %s", c_id, c.tableIdAC);
    }

    jpg.sos.spectralSelectionStart = f.readNumeric!ubyte;
    jpg.sos.spectralSelectionEnd = f.readNumeric!ubyte;
    ubyte bite = f.readNumeric!ubyte;
    jpg.sos.successiveApproximationBitHigh = bite.hiNibble;
    jpg.sos.successiveApproximationBitLow = bite.loNibble;

    writefln("SOS spectral selection start: %s", jpg.sos.spectralSelectionStart);
    writefln("SOS spectral selection end: %s", jpg.sos.spectralSelectionEnd);
    writefln("SOS successive approximation bit: %s", jpg.sos.successiveApproximationBitHigh);
    writefln("SOS successive approximation bit low: %s", jpg.sos.successiveApproximationBitLow);

    uint bytesRead = 0;
    ubyte prev = 0x00;
    bool endMarkerFound = false;
    while (!f.eof && !endMarkerFound)
    {
        bite = f.readNumeric!ubyte;
        bytesRead++;

        endMarkerFound = (prev == 0xFF && bite == 0xD9);

        if (!endMarkerFound)
        {
            if (bite != 0xFF)
            {
                ubyte datum = bite;

                if (prev == 0xFF)
                {
                    if (bite == 0x00)
                        datum = 0xFF;
                    else
                    {
                        writefln("Found marker: %X, %X", prev, bite);
                        prev = 0x00;
                        continue;
                    }
                }

                jpg.sos.encodedData ~= datum;
            }

            prev = bite;
        }
    }

    writefln("Bytes read: %s", bytesRead);
    writefln("SOS encoded data length: %s", jpg.sos.encodedData.length);
}

void decodeScanData(JPEGImage jpg)
{
    assert(jpg.sos.encodedData.length);

    //foreach(b; jpg.sos.encodedData)
    //    writef("%X ", b);
    //write("\n");

    //foreach(ref c; jpg.sos.components)
    //{
        auto c = &jpg.sos.components[0];
        auto tableDC = jpg.getHuffmanTable(0, c.tableIdDC);
        auto tableAC = jpg.getHuffmanTable(1, c.tableIdAC);

        assert(tableDC !is null);
        assert(tableAC !is null);

        uint bytePos = 0;
        uint bitPos = 0;

        ubyte decodeByte(HuffmanTreeNode* node)
        {
            assert(node !is null);
            while(!node.isLeaf)
            {
                bool bit = getBit(jpg.sos.encodedData[bytePos], 7-bitPos);
                bitPos++;
                if (bitPos == 8)
                {
                    bitPos = 0;
                    bytePos++;
                }

                if (bit) write("1");
                else write("0");

                if (bit)
                    node = node.right;
                else
                    node = node.left;
            }
            writeln("");
            assert(node !is null);
            return node.ch;
        }

        int decodeValue(ubyte len)
        {
            uint v = 0;

            uint i = 0;

            uint by = 0;
            uint bi = 0;

            while (i < len)
            {
                bool bit = getBit(jpg.sos.encodedData[bytePos], 7-bitPos);

                // FIXME: fix data ordering
                v = setBit(v, (by * 8 + bi), bit);

                bi++;
                if (bi == 8)
                {
                    bi = 0;
                    by++;
                }

                i++;

                bitPos++;
                if (bitPos == 8)
                {
                    bitPos = 0;
                    bytePos++;
                }

                if (bit) write("1");
                else write("0");
            }
            writeln("");

            bool sign = getBit(v, 0);

            write("v:  ");
            foreach(j; 0..uint.sizeof * 8)
            {
                bool bit = getBit(v, j);
                if (bit) write("1");
                else write("0");
            }
            writeln("");

            auto offset = 2^^len;

            writefln("Sign: %s", sign);
            writefln("Offset: %s", offset);
            writefln("v: %s", v);

            if (sign)
                return v;
            else
                return (-(offset-1) + v);
        }

        int[8*8] block;

        // Read DC coefficient
        ubyte dcCoefLen = decodeByte(tableDC.huffmanTree);
        writefln("DC coefficient length in bits: %s", dcCoefLen);

        int dc = 0;
        if (dcCoefLen > 0)
            dc = decodeValue(dcCoefLen);

        writefln("DC coefficient: %s", dc);
        block[0] = dc;

        // Read AC coefficients
        /*
        uint i = 1;
        bool eob = false;
        while(!eob && i < 64)
        {
            ubyte code = decodeByte(tableAC.huffmanTree);
            writefln("AC code: %X", code);

            if (code == 0x00)
                eob = true;
            else
            {
                //int ac = decodeValue(acCoefLen);
                //block[i] = ac;
                //writefln("AC coefficient: %s", ac);
                ubyte hi = hiNibble(code);
                ubyte lo = loNibble(code);
                writefln("AC code hi nibble: %X", hi);
                writefln("AC code lo nibble: %X", lo);

                uint zeroes = hi;
                foreach(j; 0..zeroes)
                {
                    block[i] = 0;
                    i++;
                }

                int ac = 0;
                if (lo > 0)
                    ac = decodeValue(lo);

                block[i] = ac;
                writefln("AC coefficient: %s", ac);
                i++;
            }
        }
        */
    //}
}