#include "pch.h"
#include "abr.h"
#include "log.h"

bool ABR::section_desc()
{
    auto sz = align4(ru32());
    auto n1 = ri32();             // some integer
    auto s = rwstring();          // maybe a string
    auto null = rkey_or_string();
    auto null_val = ri32();       // integer following the null

    // presets list
    auto name = rkey_or_string(); // "Brsh"
    auto list = rstring(4);       // "VlLs"
    if (auto presets = std::dynamic_pointer_cast<List>(call(list)))
    {
        for (auto const& pr : presets->items)
        {
            if (auto desc = std::dynamic_pointer_cast<Descriptor>(pr))
                m_presets.push_back(desc);
        }
        auto out = presets->str(0, "");
        auto lines = split(out, '\n');
        for (const auto& l : lines)
        {
            LOG("%s", l.c_str());
        }
    }
    snap();
    return true;
}

bool ABR::section_samp()
{
    auto section_size = ru32();
    auto start = pos();
    auto end = start + section_size;
    while (pos() < end)
    {
        auto samp_size = ru32();
        auto uid = rpascal();
        //printf("sample brush %s\n", uid.c_str());
        skip(4); // unknown bytes usually 00 01 00 00
        auto vm = parse_vmem();
        if (vm)
        {
            if (auto img = vm->image(true, true))
            {
                // TODO: check if uid already exists in map
                m_samples[uid] = img;
            }
        }
    }
    return true;
}

bool ABR::section_patt()
{
    auto section_size = ru32();
    auto start = pos();
    auto end = start + section_size;
    while (pos() < end)
    {
        auto patt_length = ru32(); // length of this pattern
        auto patt_version = ru32(); // = 1
        // Bitmap = 0; Grayscale = 1; Indexed = 2; RGB = 3;
        // CMYK = 4; Multichannel = 7; Duotone = 8; Lab = 9
        auto image_mode = ru32();
        if (!(image_mode == 1 || image_mode == 3))
        {
            LOG("PATT: skip image mode  %d\n", image_mode);
            skip(patt_length - 8);
            snap();
            continue;
        }
        auto point = rpoint();
        auto name = rwstring();
        auto uid = rpascal();

        // Index color table (256 * 3 RGB values): 
        // only present when image mode is indexed color
        // read(...);
        // no worries indexed is skipped anyway

        auto vm = parse_vmem();
        if (vm)
        {
            int nc = std::min((int)vm->channels.size(), 3);
            if (nc != image_mode)
            {
                LOG("PATT: image_mode (%d) and number of channels (%d) not matching\n",
                    image_mode, vm->channels.size());
            }
            if (auto img = vm->image(true, false))
            {
                // TODO: check if uid already exists in map
                m_patterns[uid] = img;
            }
        }
    }
    return true;
}

std::vector<std::shared_ptr<Brush>> ABR::compute_brushes(const std::string& path)
{
    std::vector<std::shared_ptr<Brush>> ret;
    for (auto const& p : m_presets)
    {
        auto samp = p->get<Descriptor>("Brsh");
        if (samp->class_id != "sampledBrush" && samp->class_id != "computedBrush")
        {
            LOG("unsupported brush type %s", samp->class_id.c_str());
            continue;
        }
        auto b = std::make_shared<Brush>();
        b->m_name = wstr2str(p->value<String>("Nm  "));

        // default values
        b->m_tip_color = { 0, 0, 0, 1 };
        b->m_tip_flow = .90f;
        b->m_tip_opacity = 1.f;

        b->m_tip_wet = p->value<UnitFloat>("Wtdg");
        b->m_tip_noise = p->value<Boolean>("Nose") ? 1.f : 0.f;

        b->m_tip_aspect = (1.f - samp->value<UnitFloat>("Rndn") * 0.01) * 0.5f + 0.5f;
        b->m_tip_size = samp->value<UnitFloat>("Dmtr");
        b->m_tip_spacing = samp->value<UnitFloat>("Spcn") * 0.01f;

        float tip_angle = -samp->value<UnitFloat>("Angl") / 360.f;
        b->m_tip_angle = tip_angle >= 0.f ? tip_angle : tip_angle + 1.f; // [-180, 180] -> [0, 1]

        b->m_tip_flipx = samp->value<Boolean>("flipX");
        b->m_tip_flipy = samp->value<Boolean>("flipY");

        // brush sample
        if (samp->class_id == "sampledBrush")
        {
            auto tip_uid = wstr2str(samp->value<String>("sampledData"));
            b->m_brush_path = path + "/brushes/" + tip_uid + ".png";
            b->m_brush_thumb_path = path + "/brushes/thumbs/" + tip_uid + ".png";

            const auto& samp_img = m_samples[tip_uid];
            if (samp_img->height > samp_img->width)
                b->m_tip_scale = { (float)samp_img->width / (float)samp_img->height, 1.f };
            else
                b->m_tip_scale = { 1.f, (float)samp_img->height / (float)samp_img->width };
        }
        else if (samp->class_id == "computedBrush")
        {
            if (samp->value<UnitFloat>("Hrdn") > 50.f)
            {
                b->m_brush_path = "data/brushes/Round-Hard.png";
                b->m_brush_thumb_path = "data/brushes/thumbs/Round-Hard.png";
            }
            else
            {
                b->m_brush_path = "data/brushes/Round-Brush.png";
                b->m_brush_thumb_path = "data/brushes/thumbs/Round-Brush.png";
            }
        }

        // Shape dynamics
        if (p->value<Boolean>("useTipDynamics"))
        {
            // other properties 'brushProjection', 'minimumRoundness', 'tiltScale'

            if (auto jitter_size = p->get<Descriptor>("szVr"))
            {
                b->m_jitter_scale = jitter_size->value<UnitFloat>("jitter") * 0.01f;
                // TODO: p->value<UnitFloat>("minimumDiameter") * 0.001f; // minimum size
                if (jitter_size->value<Integer>("bVTy") == 2)
                    b->m_tip_size_pressure = true;
            }

            if (auto jitter_angle = p->get<Descriptor>("angleDynamics"))
            {
                auto mode = jitter_angle->value<Integer>("bVTy");
                if (mode == 0)
                {
                    b->m_jitter_angle = jitter_angle->value<UnitFloat>("jitter") * 0.01f;
                }
                else if (mode == 5)
                {
                    b->m_jitter_angle = jitter_angle->value<UnitFloat>("jitter") * 0.01f;
                    b->m_tip_angle_init = true;
                }
                else if (mode == 6)
                {
                    b->m_jitter_angle = jitter_angle->value<UnitFloat>("jitter") * 0.01f;
                    b->m_tip_angle_follow = true;
                }
            }

            if (auto roundness = p->get<Descriptor>("roundnessDynamics"))
            {
                b->m_jitter_aspect = roundness->value<UnitFloat>("jitter") * 0.01;
            }

            b->m_tip_randflipx = p->value<Boolean>("flipX");
            b->m_tip_randflipy = p->value<Boolean>("flipY");
        }

        // Transfer settings
        if (p->value<Boolean>("usePaintDynamics"))
        {
            auto jitter_opacity = p->get<Descriptor>("opVr");
            if (jitter_opacity)
            {
                b->m_jitter_opacity = jitter_opacity->value<UnitFloat>("jitter") * 0.01f;
                // TODO: jitter_opacity->value<UnitFloat>("Mnm ") * 0.01f; // minimum size
                if (jitter_opacity->value<Integer>("bVTy") == 2)
                    b->m_tip_opacity_pressure = true;
            }
            
            auto jitter_flow = p->get<Descriptor>("prVr");
            if (jitter_flow)
            {
                b->m_jitter_flow = jitter_flow->value<UnitFloat>("jitter") * 0.01f;
                // TODO: m_jitter_flow->value<UnitFloat>("Mnm ") * 0.01f; // minimum size
                if (jitter_flow->value<Integer>("bVTy") == 2)
                    b->m_tip_flow_pressure = true;
            }
            
        }

        std::vector<std::string> modes = {
            "normal",       // normal (not in Photoshop)
            "Mltp",         // multiply
            "Sbtr",         // subtract
            "Drkn",         // darken
            "Ovrl",         // overlay
            "CDdg",         // color dodge
            "CBrn",         // color burn
            "linearBurn",   // linear burn
            "hardMix",      // hard mix
            "linearHeight", // linear height
            "Hght",         // height
        };

        // pattern
        if (auto patt = p->get<Descriptor>("Txtr"))
        {
            auto patt_uid = wstr2str(patt->value<String>("Idnt"));
            b->m_pattern_path = path + "/patterns/" + patt_uid + ".png";
            b->m_pattern_thumb_path = path + "/patterns/thumbs/" + patt_uid + ".png";
            
            b->m_pattern_invert = p->value<Boolean>("InvT");
            b->m_pattern_eachsample = p->value<Boolean>("TxtC");
            b->m_pattern_depth = p->value<UnitFloat>("textureDepth") * 0.01f; // [0, 100] -> [0, 1]
            b->m_pattern_scale = p->value<UnitFloat>("textureScale") * 0.01f; // [0, 1000] -> [0, 1]
            b->m_pattern_brightness = (float)(p->value<Integer>("textureBrightness") + 150) / 300.f; // [-150, 150] -> [0, 1]

            int raw_contrast = p->value<Integer>("textureContrast"); // [-50, 0, 100] -> [0, 1]
            b->m_pattern_contrast = raw_contrast / (raw_contrast < 0 ? 100.f : 200.f) + 0.5f;

            // blending mode
            std::string blend_mode = p->value<Enum>("textureBlendMode");
            auto bm_it = std::find(modes.begin(), modes.end(), blend_mode);
            if (bm_it != modes.end())
                b->m_pattern_blend_mode = std::distance(modes.begin(), bm_it);

            b->m_pattern_enabled = p->value<Boolean>("useTexture");
        }

        // dual brush
        auto db = p->get<Descriptor>("dualBrush");
        if (db && db->value<Boolean>("useDualBrush"))
        {
            auto samp = db->get<Descriptor>("Brsh");
            if (samp->class_id != "sampledBrush" && samp->class_id != "computedBrush")
            {
                LOG("unsupported brush type %s", samp->class_id.c_str());
                continue;
            }
            //b->m_name = wstr2str(p->value<String>("Nm  "));

            // default values
            //b->m_tip_color = { 0, 0, 0, 1 };
            b->m_dual_flow = .90f;
            b->m_dual_opacity = 1.f;

            b->m_dual_aspect = (1.f - samp->value<UnitFloat>("Rndn") * 0.01) * 0.5f + 0.5f;
            b->m_dual_size = samp->value<UnitFloat>("Dmtr") / b->m_tip_size;
            b->m_dual_spacing = samp->value<UnitFloat>("Spcn") * 0.01f;

            float tip_angle = -samp->value<UnitFloat>("Angl") / 360.f;
            b->m_dual_angle = tip_angle >= 0.f ? tip_angle : tip_angle + 1.f; // [-180, 180] -> [0, 1]
            
            //b->m_tip_wet = p->value<UnitFloat>("Wtdg");
            //b->m_tip_noise = (float)samp->value<UnitFloat>("Nose");
            b->m_tip_flipx = samp->value<Boolean>("flipX");
            b->m_tip_flipy = samp->value<Boolean>("flipY");
            
            b->m_dual_randflip = db->value<Boolean>("Flip");
            b->m_dual_scatter_bothaxis = db->value<Boolean>("bothAxes");

            if (db->value<Boolean>("useScatter"))
            {
                auto scatter = db->get<Descriptor>("scatterDynamics");
                b->m_dual_scatter = scatter->value<UnitFloat>("jitter") * 0.01f;
            }

            // brush sample
            if (samp->class_id == "sampledBrush")
            {
                auto tip_uid = wstr2str(samp->value<String>("sampledData"));
                b->m_dual_path = path + "/brushes/" + tip_uid + ".png";
                b->m_dual_thumb_path = path + "/brushes/thumbs/" + tip_uid + ".png";

                const auto& samp_img = m_samples[tip_uid];
                //b->m_tip_width = (float)samp_img->width / (float)samp_img->height;
            }
            else if (samp->class_id == "computedBrush")
            {
                if (samp->value<UnitFloat>("Hrdn") > 50.f)
                {
                    b->m_dual_path = "data/brushes/Round-Hard.png";
                    b->m_dual_thumb_path = "data/brushes/thumbs/Round-Hard.png";
                }
                else
                {
                    b->m_dual_path = "data/brushes/Round-Brush.png";
                    b->m_dual_thumb_path = "data/brushes/thumbs/Round-Brush.png";
                }
            }

            // blending mode
            std::string blend_mode = db->value<Enum>("BlnM");
            auto bm_it = std::find(modes.begin(), modes.end(), blend_mode);
            if (bm_it != modes.end())
                b->m_dual_blend_mode = std::distance(modes.begin(), bm_it);

            b->m_dual_enabled = db->value<Boolean>("useDualBrush");
        }

        ret.push_back(b);
    }
    return ret;
}

std::shared_ptr<ABR::VMArray> ABR::parse_vmem()
{
    // Virtual Memory Array List
    auto vmem_version = ru32(); // = 3
    assert(vmem_version == 3);
    auto vmem_length = ru32();
    // TODO: check if at the end there's good data
    // check if the bounds are within the parent's size
    auto vmem_rect = rrect();
    auto vmem_channels = ru32();
    // The following is a virtual memory array,
    // repeated for the number of channels 
    // + one for a user mask + one for a sheet mask.
    vmem_channels += 2; // user and sheet mask
    auto ret = std::make_shared<VMArray>(vmem_version, vmem_rect);
    for (int ch = 0; ch < vmem_channels; ch++)
    {
        auto array_written = ru32(); // skip if 0
        if (array_written == 0)
            continue;
        auto length = ru32(); // skip if 0, length is from the next field to the end
        if (length == 0)
            continue;
        auto depth = ru32(); // Pixel depth: 1, 8, 16 or 32
        auto rect = rrect();
        auto depth2 = ru16(); // again?
        auto compression = ru8(); // 1 = zip
        if (depth != 8)
        {
            LOG("unsupported depth %d bits\n", depth);
            skip(length - 23);
            continue;
        }
        if (compression == 0)
        {
            int data_size = (depth >> 3) * rect.area();
            ret->channels.emplace_back(depth, rect, compression, rraw(data_size));
        }
        else if (compression == 1)
        {
            auto start = pos();
            auto height = rect.height();
            // contain the compressed length of each scanline
            std::vector<uint16_t> scanlines;
            scanlines.reserve(height);
            for (int i = 0; i < height; i++)
                scanlines.push_back(ru16());
            std::vector<uint8_t> raw;
            for (auto sl : scanlines)
            {
                auto decoded = rrle(sl);
                raw.insert(raw.end(), decoded.begin(), decoded.end());
            }
            ret->channels.emplace_back(depth, rect, compression, raw);
            auto len = pos() - start;
        }
        else
        {
            LOG("unsupported compression mode %d\n", compression);
            skip(length - 23);
            continue;
        }
    }
    snap();
    return ret;
}

std::shared_ptr<ABR::List> ABR::parse_vlls()
{
    auto ret = std::make_shared<List>();
    auto count = ru32();
    //printf("list: %d\n", count);
    for (int i = 0; i < count; i++)
    {
        auto type = rstring(4);
        auto item = call(type);
        if (!item)
            return nullptr;
        ret->items.push_back(item);
    }
    return ret;
}

std::shared_ptr<ABR::String> ABR::parse_text()
{
    auto ret = std::make_shared<String>();
    ret->value = rwstring();
    //wprintf(L"text: %s\n", ret->value.c_str());
    return ret;
}

std::shared_ptr<ABR::Descriptor> ABR::parse_objc()
{
    auto ret = std::make_shared<Descriptor>();
    ret->name = rwstring();
    ret->class_id = rkey_or_string();
    auto count = ru32();
    //printf("objc type %s, %d props\n", ret->class_id.c_str(), count);
    for (int i = 0; i < count; i++)
    {
        auto key = rkey_or_string();
        auto type = rstring(4);
        //printf("prop %s\n", t.c_str());
        auto property = call(type);
        if (!property)
            return nullptr;
        if (ret->props.find(key) != ret->props.end())
            LOG("DUPLICATE prop %s\n", key.c_str());
        ret->props[key] = property;
    }
    return ret;
}

std::shared_ptr<ABR::UnitFloat> ABR::parse_untf()
{
    auto ret = std::make_shared<UnitFloat>();
    ret->unit = rstring(4);
    ret->value = rdbl();
    //printf("float %s: %f\n", ret->unit.c_str(), ret->value);
    return ret;
}

std::shared_ptr<ABR::Boolean> ABR::parse_bool()
{
    auto ret = std::make_shared<Boolean>();
    ret->value = ru8();
    //printf("bool: %s\n", ret->value ? "true" : "false");
    return ret;
}

std::shared_ptr<ABR::Integer> ABR::parse_long()
{
    auto ret = std::make_shared<Integer>();
    ret->value = ru32();
    //printf("long: %d\n", ret->value);
    return ret;
}

std::shared_ptr<ABR::Double> ABR::parse_doub()
{
    auto ret = std::make_shared<Double>();
    ret->value = rdbl();
    //printf("double: %d\n", ret->value);
    return ret;
}

std::shared_ptr<ABR::Enum> ABR::parse_enum()
{
    auto ret = std::make_shared<Enum>();
    ret->type = rkey_or_string();
    ret->value = rkey_or_string();
    //printf("enum: %s %s\n", t.c_str(), e.c_str());
    return ret;
}

std::shared_ptr<ABR::RawData> ABR::parse_tdta()
{
    auto ret = std::make_shared<RawData>();
    ret->data = rraw();
    return ret;
}

ABR::ABR()
{
    m_parser_table = std::map<std::string, std::function<ABR::Type::Ref()>>
    {
        { "VlLs", std::bind(&ABR::parse_vlls, this) },
        { "TEXT", std::bind(&ABR::parse_text, this) },
        { "Objc", std::bind(&ABR::parse_objc, this) },
        { "UntF", std::bind(&ABR::parse_untf, this) },
        { "bool", std::bind(&ABR::parse_bool, this) },
        { "long", std::bind(&ABR::parse_long, this) },
        { "doub", std::bind(&ABR::parse_doub, this) },
        { "enum", std::bind(&ABR::parse_enum, this) },
        { "tdta", std::bind(&ABR::parse_tdta, this) },
    };
}

bool ABR::open(const std::string& path)
{
    Asset asset;
    if (asset.open(path.c_str()))
    {
        init(asset.read_all(), asset.m_len, BinaryStream::ByteOrder::BigEndian);
        auto version_major = ru16();
        auto version_minor = ru16();
        LOG("ABR %d.%d\n", version_major, version_minor);
        while (!eof())
        {
            if (rstring(4) != "8BIM")
                return false;
            auto t = rstring(4);
            if (t == "desc")
            {
                section_desc();
            }
            else if (t == "patt")
            {
                section_patt();
            }
            else if (t == "samp")
            {
                section_samp();
            }
            else
            {
                LOG("skip section %s\n", t.c_str());
                skip(align4(ru32()));
            }
        }
        return true;
    }
    return false;
}