#!/usr/bin/env python3

# Copyright (c) 2023, Lucas Chollet <lucas.chollet@serenityos.org>
#
# SPDX-License-Identifier: BSD-2-Clause

import argparse
import re

from collections import namedtuple
from enum import Enum
from pathlib import Path
from typing import Any
from typing import List
from typing import Type


class EnumWithExportName(Enum):
    @classmethod
    def export_name(cls) -> str:
        return cls.__name__


class TIFFType(EnumWithExportName):
    @classmethod
    def export_name(cls) -> str:
        return "Type"

    def __new__(cls, *args):
        obj = object.__new__(cls)
        obj._value_ = args[0]
        obj.size = args[1]
        return obj

    # First value is the underlying u16, second one is the size in bytes
    Byte = 1, 1
    ASCII = 2, 1
    UnsignedShort = 3, 2
    UnsignedLong = 4, 4
    UnsignedRational = 5, 8
    Undefined = 7, 1
    SignedLong = 9, 4
    SignedRational = 10, 8
    Float = 11, 4
    Double = 12, 8
    IFD = 13, 4
    UTF8 = 129, 1


class Predictor(EnumWithExportName):
    NoPrediction = 1
    HorizontalDifferencing = 2


class Compression(EnumWithExportName):
    NoCompression = 1
    CCITTRLE = 2
    Group3Fax = 3
    Group4Fax = 4
    LZW = 5
    OldJPEG = 6
    JPEG = 7
    AdobeDeflate = 8
    PackBits = 32773
    PixarDeflate = 32946  # This is the old (and deprecated) code for AdobeDeflate


class PhotometricInterpretation(EnumWithExportName):
    WhiteIsZero = 0
    BlackIsZero = 1
    RGB = 2
    RGBPalette = 3
    TransparencyMask = 4
    CMYK = 5
    YCbCr = 6
    CIELab = 8


class FillOrder(EnumWithExportName):
    LeftToRight = 1
    RightToLeft = 2


class Orientation(EnumWithExportName):
    Default = 1
    FlipHorizontally = 2
    Rotate180 = 3
    FlipVertically = 4
    Rotate90ClockwiseThenFlipHorizontally = 5
    Rotate90Clockwise = 6
    FlipHorizontallyThenRotate90Clockwise = 7
    Rotate90CounterClockwise = 8


class PlanarConfiguration(EnumWithExportName):
    Chunky = 1
    Planar = 2


class ResolutionUnit(EnumWithExportName):
    NoAbsolute = 1
    Inch = 2
    Centimeter = 3


class SampleFormat(EnumWithExportName):
    Unsigned = 1
    Signed = 2
    Float = 3
    Undefined = 4


class ExtraSample(EnumWithExportName):
    Unspecified = 0
    AssociatedAlpha = 1
    UnassociatedAlpha = 2


tag_fields = ["id", "types", "counts", "default", "name", "associated_enum", "is_required"]

Tag = namedtuple(
    "Tag",
    field_names=tag_fields,
    defaults=(None,) * len(tag_fields),
)

# FIXME: Some tag have only a few allowed values, we should ensure that
known_tags: List[Tag] = [
    Tag("256", [TIFFType.UnsignedShort, TIFFType.UnsignedLong], [1], None, "ImageWidth", is_required=True),
    Tag("257", [TIFFType.UnsignedShort, TIFFType.UnsignedLong], [1], None, "ImageLength", is_required=True),
    Tag("258", [TIFFType.UnsignedShort], [], None, "BitsPerSample", is_required=False),
    Tag("259", [TIFFType.UnsignedShort], [1], None, "Compression", Compression, is_required=True),
    Tag(
        "262",
        [TIFFType.UnsignedShort],
        [1],
        None,
        "PhotometricInterpretation",
        PhotometricInterpretation,
        is_required=True,
    ),
    Tag("266", [TIFFType.UnsignedShort], [1], FillOrder.LeftToRight, "FillOrder", FillOrder),
    Tag("271", [TIFFType.ASCII], [], None, "Make"),
    Tag("272", [TIFFType.ASCII], [], None, "Model"),
    Tag("273", [TIFFType.UnsignedShort, TIFFType.UnsignedLong], [], None, "StripOffsets", is_required=False),
    Tag("274", [TIFFType.UnsignedShort], [1], Orientation.Default, "Orientation", Orientation),
    Tag("277", [TIFFType.UnsignedShort], [1], None, "SamplesPerPixel", is_required=False),
    Tag("278", [TIFFType.UnsignedShort, TIFFType.UnsignedLong], [1], None, "RowsPerStrip", is_required=False),
    Tag("279", [TIFFType.UnsignedShort, TIFFType.UnsignedLong], [], None, "StripByteCounts", is_required=False),
    Tag("282", [TIFFType.UnsignedRational], [1], None, "XResolution"),
    Tag("283", [TIFFType.UnsignedRational], [1], None, "YResolution"),
    Tag("284", [TIFFType.UnsignedShort], [1], PlanarConfiguration.Chunky, "PlanarConfiguration", PlanarConfiguration),
    Tag("285", [TIFFType.ASCII], [], None, "PageName"),
    Tag("292", [TIFFType.UnsignedLong], [1], 0, "T4Options"),
    Tag("296", [TIFFType.UnsignedShort], [1], ResolutionUnit.Inch, "ResolutionUnit", ResolutionUnit),
    Tag("305", [TIFFType.ASCII], [], None, "Software"),
    Tag("306", [TIFFType.ASCII], [20], None, "DateTime"),
    Tag("315", [TIFFType.ASCII], [], None, "Artist"),
    Tag("317", [TIFFType.UnsignedShort], [1], Predictor.NoPrediction, "Predictor", Predictor),
    Tag("320", [TIFFType.UnsignedShort], [], None, "ColorMap"),
    Tag("322", [TIFFType.UnsignedShort, TIFFType.UnsignedLong], [1], None, "TileWidth"),
    Tag("323", [TIFFType.UnsignedShort, TIFFType.UnsignedLong], [1], None, "TileLength"),
    Tag("324", [TIFFType.UnsignedShort, TIFFType.UnsignedLong], [], None, "TileOffsets"),
    Tag("325", [TIFFType.UnsignedShort, TIFFType.UnsignedLong], [], None, "TileByteCounts"),
    Tag("338", [TIFFType.UnsignedShort], [], None, "ExtraSamples", ExtraSample),
    Tag("339", [TIFFType.UnsignedShort], [], SampleFormat.Unsigned, "SampleFormat", SampleFormat),
    Tag("34665", [TIFFType.UnsignedLong, TIFFType.IFD], [1], None, "ExifIFD"),
    Tag("34675", [TIFFType.Undefined], [], None, "ICCProfile"),
]

HANDLE_TAG_SIGNATURE_TEMPLATE = (
    "ErrorOr<void> {namespace}handle_tag(Function<ErrorOr<void>(u32)>&& subifd_handler, "
    "ExifMetadata& metadata, u16 tag, {namespace}Type type, u32 count, "
    "Vector<{namespace}Value>&& value)"
)
HANDLE_TAG_SIGNATURE = HANDLE_TAG_SIGNATURE_TEMPLATE.format(namespace="")
HANDLE_TAG_SIGNATURE_TIFF_NAMESPACE = HANDLE_TAG_SIGNATURE_TEMPLATE.format(namespace="TIFF::")

ENSURE_BASELINE_TAG_PRESENCE = "ErrorOr<void> ensure_baseline_tags_are_present(ExifMetadata const& metadata)"
TIFF_TYPE_FROM_U16 = "ErrorOr<Type> tiff_type_from_u16(u16 type)"
SIZE_OF_TIFF_TYPE = "u8 size_of_type(Type type)"

LICENSE = R"""/*
 * Copyright (c) 2023, Lucas Chollet <lucas.chollet@serenityos.org>
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */"""


def export_enum_to_cpp(e: Type[EnumWithExportName]) -> str:
    output = f"enum class {e.export_name()} {{\n"

    for entry in e:
        output += f"    {entry.name} = {entry.value},\n"

    output += "};\n"
    return output


def export_enum_to_string_converter(enums: List[Type[EnumWithExportName]]) -> str:
    stringifier_internals = []
    for e in enums:
        single_stringifier = Rf"""    if constexpr (IsSame<E, {e.export_name()}>) {{
        switch (value) {{
            default:
                return "Invalid value for {e.export_name()}"sv;"""
        for entry in e:
            single_stringifier += Rf"""
            case {e.export_name()}::{entry.name}:
                return "{entry.name}"sv;"""

        single_stringifier += R"""
        }
    }"""
        stringifier_internals.append(single_stringifier)

    stringifier_internals_str = "\n".join(stringifier_internals)

    out = Rf"""template<Enum E>
StringView name_for_enum_tag_value(E value) {{
{stringifier_internals_str}
    VERIFY_NOT_REACHED();
}}"""

    return out


def export_tag_related_enums(tags: List[Tag]) -> str:
    exported_enums = []
    for tag in tags:
        if tag.associated_enum:
            exported_enums.append(export_enum_to_cpp(tag.associated_enum))

    return "\n".join(exported_enums)


def promote_type(t: TIFFType) -> TIFFType:
    if t == TIFFType.UnsignedShort:
        return TIFFType.UnsignedLong
    if t == TIFFType.Float:
        return TIFFType.Double
    return t


def tiff_type_to_cpp(t: TIFFType, with_promotion: bool = True) -> str:
    # To simplify the code generator and the ExifMetadata class API, all u16 are promoted to u32
    # Note that the Value<> type doesn't include u16 for this reason
    if with_promotion:
        t = promote_type(t)
    if t in [TIFFType.ASCII, TIFFType.UTF8]:
        return "String"
    if t == TIFFType.Undefined:
        return "ByteBuffer"
    if t == TIFFType.UnsignedShort:
        return "u16"
    if t == TIFFType.UnsignedLong or t == TIFFType.IFD:
        return "u32"
    if t == TIFFType.UnsignedRational:
        return "TIFF::Rational<u32>"
    if t == TIFFType.Float:
        return "float"
    if t == TIFFType.Double:
        return "double"
    raise RuntimeError(f'Type "{t}" not recognized, please update tiff_type_to_read_only_cpp()')


def is_container(t: TIFFType) -> bool:
    """
    Some TIFF types are defined on the unit scale but are intended to be used within a collection.
    An example of that are ASCII strings defined as N * byte. Let's intercept that and generate
    a nice API instead of Vector<u8>.
    """
    return t in [TIFFType.ASCII, TIFFType.Byte, TIFFType.Undefined, TIFFType.UTF8]


def export_promoter() -> str:
    output = R"""template<typename T>
struct TypePromoter {
    using Type = T;
};
"""
    specialization_template = R"""template<>
struct TypePromoter<{}> {{
    using Type = {};
}};
"""
    for t in TIFFType:
        if promote_type(t) != t:
            output += specialization_template.format(tiff_type_to_cpp(t, with_promotion=False), tiff_type_to_cpp(t))

    return output


def retrieve_biggest_type(types: List[TIFFType]) -> TIFFType:
    return TIFFType(max([t.value for t in types]))


def pascal_case_to_snake_case(name: str) -> str:
    name = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", name)
    return re.sub("([a-z0-9])([A-Z])", r"\1_\2", name).lower()


def default_value_to_cpp(value: Any) -> str:
    if isinstance(value, EnumWithExportName):
        return f"TIFF::{value.export_name()}::{value.name}"
    return str(value)


def generate_getter(tag: Tag) -> str:
    biggest_type = retrieve_biggest_type(tag.types)
    variant_inner_type = tiff_type_to_cpp(biggest_type)

    extracted_value_template = f"(*possible_value)[{{}}].get<{variant_inner_type}>()"

    tag_final_type = variant_inner_type
    if tag.associated_enum:
        tag_final_type = f"TIFF::{tag.associated_enum.__name__}"
        extracted_value_template = f"static_cast<{tag_final_type}>({extracted_value_template})"

    single_count = len(tag.counts) == 1 and tag.counts[0] == 1 or is_container(biggest_type)
    if single_count:
        return_type = tag_final_type
        if is_container(biggest_type):
            return_type += " const&"
        unpacked_if_needed = f"return {extracted_value_template.format(0)};"
    else:
        if len(tag.counts) == 1:
            container_type = f"Array<{tag_final_type}, {tag.counts[0]}>"
            container_initialization = f"{container_type} tmp{{}};"
        else:
            container_type = f"Vector<{tag_final_type}>"
            container_initialization = Rf"""{container_type} tmp{{}};
        auto maybe_failure = tmp.try_resize(possible_value->size());
        if (maybe_failure.is_error())
            return OptionalNone {{}};
        """

        return_type = container_type
        unpacked_if_needed = Rf"""
        {container_initialization}
        for (u32 i = 0; i < possible_value->size(); ++i)
            tmp[i] = {extracted_value_template.format("i")};

        return tmp;"""

    signature = Rf"    Optional<{return_type}> {pascal_case_to_snake_case(tag.name)}() const"

    if tag.default is not None and single_count:
        return_if_empty = f"{default_value_to_cpp(tag.default)}"
    else:
        return_if_empty = "OptionalNone {}"

    body = Rf"""
    {{
        auto const& possible_value = m_data.get("{tag.name}"sv);
        if (!possible_value.has_value())
            return {return_if_empty};
        {unpacked_if_needed}
    }}
"""

    return signature + body


def generate_metadata_class(tags: List[Tag]) -> str:
    getters = "\n".join([generate_getter(tag) for tag in tags])

    output = Rf"""class ExifMetadata : public Metadata {{
public:
    virtual ~ExifMetadata() = default;

{getters}
private:
    friend {HANDLE_TAG_SIGNATURE_TIFF_NAMESPACE};

    virtual void fill_main_tags() const override {{
        if (model().has_value())
            m_main_tags.set("Model"sv, model().value());
        if (make().has_value())
            m_main_tags.set("Manufacturer"sv, make().value());
        if (software().has_value())
            m_main_tags.set("Software"sv, software().value());
        if (date_time().has_value())
            m_main_tags.set("Creation Time"sv, date_time().value());
        if (artist().has_value())
            m_main_tags.set("Author"sv, artist().value());
    }}

    void add_entry(StringView key, Vector<TIFF::Value>&& value) {{
        m_data.set(key, move(value));
    }}

    HashMap<StringView, Vector<TIFF::Value>> m_data;
}};
"""

    return output


def generate_metadata_file(tags: List[Tag]) -> str:
    output = Rf"""{LICENSE}

#pragma once

#include <AK/HashMap.h>
#include <AK/Variant.h>
#include <AK/Vector.h>
#include <LibGfx/Size.h>
#include <LibGfx/ImageFormats/ImageDecoder.h>

namespace Gfx {{

class ExifMetadata;

namespace TIFF {{

{export_enum_to_cpp(TIFFType)}

template<OneOf<u32, i32> x32>
struct Rational {{
    using Type = x32;
    x32 numerator;
    x32 denominator;

    double as_double() const {{
        return static_cast<double>(numerator) / denominator;
    }}
}};

{export_promoter()}

// Note that u16 is not include on purpose
using Value = Variant<ByteBuffer, String, u32, Rational<u32>, i32, Rational<i32>, double>;

{export_tag_related_enums(known_tags)}

{export_enum_to_string_converter([tag.associated_enum for tag in known_tags if tag.associated_enum] + [TIFFType])}

{HANDLE_TAG_SIGNATURE};
{ENSURE_BASELINE_TAG_PRESENCE};
{TIFF_TYPE_FROM_U16};
{SIZE_OF_TIFF_TYPE};

}}

{generate_metadata_class(tags)}

}}

template<typename T>
struct AK::Formatter<Gfx::TIFF::Rational<T>> : Formatter<FormatString> {{
    ErrorOr<void> format(FormatBuilder& builder, Gfx::TIFF::Rational<T> value)
    {{
        return Formatter<FormatString>::format(builder, "{{}} ({{}}/{{}})"sv,
            value.as_double(), value.numerator, value.denominator);
    }}
}};

template<>
struct AK::Formatter<Gfx::TIFF::Value> : Formatter<FormatString> {{
    ErrorOr<void> format(FormatBuilder& builder, Gfx::TIFF::Value const& value)
    {{
        String content;
        value.visit(
            [&](ByteBuffer const& buffer) {{
                content = MUST(String::formatted("Buffer of size: {{}}", buffer.size()));
            }},
            [&](auto const& other) {{
                content = MUST(String::formatted("{{}}", other));
            }}
        );

        return Formatter<FormatString>::format(builder, "{{}}"sv, content);
    }}
}};
"""
    return output


def generate_tag_handler(tag: Tag) -> str:
    not_in_type_list = f"({' && '.join([f'type != Type::{t.name}' for t in tag.types])})"

    not_in_count_list = ""
    if len(tag.counts) != 0:
        not_in_count_list = f"|| ({' && '.join([f'count != {c}' for c in tag.counts])})"
    pre_condition = Rf"""if ({not_in_type_list}
            {not_in_count_list})
            return Error::from_string_literal("TIFFImageDecoderPlugin: Tag {tag.name} invalid");"""

    check_value = ""
    if tag.associated_enum is not None:
        not_in_value_list = f"({' && '.join([f'v != {v.value}' for v in tag.associated_enum])})"
        check_value = Rf"""
        for (u32 i = 0; i < value.size(); ++i) {{
            TRY(value[i].visit(
                []({tiff_type_to_cpp(tag.types[0])} const& v) -> ErrorOr<void> {{
                    if ({not_in_value_list})
                        return Error::from_string_literal("TIFFImageDecoderPlugin: Invalid value for tag {tag.name}");
                    return {{}};
                }},
                [&](auto const&) -> ErrorOr<void> {{
                    VERIFY_NOT_REACHED();
                }})
            );
        }}
"""

    handle_subifd = ""
    if TIFFType.IFD in tag.types:
        if tag.counts != [1]:
            raise RuntimeError("Accessing `value[0]` in the C++ code might fail!")
        handle_subifd = f"TRY(subifd_handler(value[0].get<{tiff_type_to_cpp(TIFFType.IFD)}>()));"

    output = Rf"""    case {tag.id}:
        // {tag.name}

        dbgln_if(TIFF_DEBUG, "{tag.name}({{}}): {{}}", name_for_enum_tag_value(type), format_tiff_value(tag, value));

        {pre_condition}
        {check_value}
        {handle_subifd}
        metadata.add_entry("{tag.name}"sv, move(value));
        break;
"""

    return output


def generate_tag_handler_file(tags: List[Tag]) -> str:
    formatter_for_tag_with_enum = "\n".join(
        [
            Rf"""        case {tag.id}:
            return MUST(String::from_utf8(
                name_for_enum_tag_value(static_cast<{tag.associated_enum.export_name()}>(v.get<u32>()))));"""
            for tag in tags
            if tag.associated_enum
        ]
    )

    ensure_tags_are_present = "\n".join(
        [
            Rf"""    if (!metadata.{pascal_case_to_snake_case(tag.name)}().has_value())
        return Error::from_string_literal("Unable to decode image, missing required tag {tag.name}.");
"""
            for tag in filter(lambda tag: tag.is_required, known_tags)
        ]
    )

    tiff_type_from_u16_cases = "\n".join(
        [
            Rf"""    case to_underlying(Type::{t.name}):
        return Type::{t.name};"""
            for t in TIFFType
        ]
    )

    size_of_tiff_type_cases = "\n".join(
        [
            Rf"""    case Type::{t.name}:
        return {t.size};"""
            for t in TIFFType
        ]
    )

    output = Rf"""{LICENSE}

#include <AK/Debug.h>
#include <AK/String.h>
#include <LibGfx/ImageFormats/TIFFMetadata.h>

namespace Gfx::TIFF {{

static String value_formatter(u32 tag_id, Value const& v) {{
    switch (tag_id) {{
{formatter_for_tag_with_enum}
        default:
            return MUST(String::formatted("{{}}", v));
    }}
}}

[[maybe_unused]] static String format_tiff_value(u32 tag_id, Vector<Value> const& values) {{
    if (values.size() == 1)
        return MUST(String::formatted("{{}}", value_formatter(tag_id, values[0])));

    StringBuilder builder;
    builder.append('[');

    for (u32 i = 0; i < values.size(); ++i) {{
        builder.appendff("{{}}", value_formatter(tag_id, values[i]));
        if (i != values.size() - 1)
            builder.append(", "sv);
    }}

    builder.append(']');
    return MUST(builder.to_string());
}}

{ENSURE_BASELINE_TAG_PRESENCE}
{{
{ensure_tags_are_present}
    return {{}};
}}

{TIFF_TYPE_FROM_U16}
{{
    switch (type) {{
{tiff_type_from_u16_cases}
    default:
        return Error::from_string_literal("TIFFImageDecoderPlugin: Unknown type");
    }}
}}

{SIZE_OF_TIFF_TYPE}
{{
    switch (type) {{
{size_of_tiff_type_cases}
    default:
        VERIFY_NOT_REACHED();
    }}
}}

{HANDLE_TAG_SIGNATURE}
{{
    switch (tag) {{
"""

    output += "\n".join([generate_tag_handler(t) for t in tags])

    output += R"""
    default:
        dbgln_if(TIFF_DEBUG, "UnknownTag({}, {}): {}",
                tag, name_for_enum_tag_value(type), format_tiff_value(tag, value));
    }

    return {};
}

}
"""
    return output


def update_file(target: Path, new_content: str):
    should_update = True

    if target.exists():
        with target.open("r") as file:
            content = file.read()
            if content == new_content:
                should_update = False

    if should_update:
        with target.open("w") as file:
            file.write(new_content)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("-o", "--output")
    args = parser.parse_args()

    output_path = Path(args.output)

    update_file(output_path / "TIFFMetadata.h", generate_metadata_file(known_tags))
    update_file(output_path / "TIFFTagHandler.cpp", generate_tag_handler_file(known_tags))


if __name__ == "__main__":
    main()