# -*- coding: utf-8 -*-

import arcpy  # type: ignore


class Toolbox:
    def __init__(self):
        self.label = "Feature Renderer"
        self.alias = "Feature Renderer"

        # List of tool classes associated with this toolbox
        self.tools = [FeatureRenderer]


class FeatureRenderer:
    def __init__(self):
        self.label = "Feature Renderer"
        self.description = "Unique value renderer for geologic features"

    def getParameterInfo(self):
        # Define parameter for selecting style files
        style_files = arcpy.Parameter(
            displayName="Select Style Files",
            name="style_files",
            datatype="DEFile",
            parameterType="Required",
            direction="Input",
            multiValue=True,
        )

        # Set filter to only show .stylx files in the file picker
        style_files.filter.list = ["stylx"]

        # Define parameter for selecting feature classes or layers
        layers = arcpy.Parameter(
            displayName="Select Feature Classes or Layers",
            name="fc_or_layer",
            datatype="GPFeatureLayer",
            parameterType="Required",
            direction="Input",
            multiValue=True,
        )

        # Define parameter for selecting a table from the enterprise geodatabase
        table_param = arcpy.Parameter(
            displayName="Select Legend Table",
            name="table",
            datatype="DETable",
            parameterType="Required",
            direction="Input",
        )

        # Define parameter for selecting a primary key
        field_param_1 = arcpy.Parameter(
            displayName="Select Primary Key Field (LEG_ID)",
            name="primary_key",
            datatype="Field",
            parameterType="Required",
            direction="Input",
            multiValue=False,
        )

        field_param_1.parameterDependencies = [table_param.name]

        # Define parameter for selecting a fill symbol code field
        field_param_2 = arcpy.Parameter(
            displayName="Select Fill Symbol Code Field",
            name="fill_symbol",
            datatype="Field",
            parameterType="Required",
            direction="Input",
            multiValue=False,
        )

        field_param_2.parameterDependencies = [table_param.name]

        # Define parameter for selecting a stroke symbol code field
        field_param_3 = arcpy.Parameter(
            displayName="Select Stroke Symbol Code Field",
            name="line_symbol",
            datatype="Field",
            parameterType="Required",
            direction="Input",
            multiValue=False,
        )

        field_param_3.parameterDependencies = [table_param.name]

        # Define parameter for selecting a marker symbol code field
        field_param_4 = arcpy.Parameter(
            displayName="Select Marker Symbol Code Field",
            name="point_symbol",
            datatype="Field",
            parameterType="Required",
            direction="Input",
            multiValue=False,
        )

        field_param_4.parameterDependencies = [table_param.name]

        # Define parameter for selecting a legend text field
        field_param_5 = arcpy.Parameter(
            displayName="Select Legend Text Field 1",
            name="legend_text_1",
            datatype="Field",
            parameterType="Optional",
            direction="Input",
            multiValue=False,
        )

        field_param_5.parameterDependencies = [table_param.name]

        # Define parameter for selecting a legend text field
        field_param_6 = arcpy.Parameter(
            displayName="Select Legend Text Field 2",
            name="legend_text_2",
            datatype="Field",
            parameterType="Optional",
            direction="Input",
            multiValue=False,
        )

        field_param_6.parameterDependencies = [table_param.name]

        # Define parameter for label field
        label_field_1 = arcpy.Parameter(
            name="label_field_1",
            displayName="Select Label Field 1",
            datatype="Field",
            direction="Input",
            parameterType="Optional",
            category="Define Label",
        )

        # Define parameter for delimiter
        label_delimiter = arcpy.Parameter(
            name="label_delimiter",
            displayName="Delimiter",
            datatype="GPString",
            direction="Input",
            parameterType="Optional",
            category="Define Label",
        )

        label_delimiter.value = "-"

        # Define parameter for label field
        label_field_2 = arcpy.Parameter(
            name="label_field_2",
            displayName="Select Label Field 2",
            datatype="Field",
            direction="Input",
            parameterType="Optional",
            category="Define Label",
        )

        label_field_1.parameterDependencies = [table_param.name]
        label_field_2.parameterDependencies = [table_param.name]

        # Define parameter for heading field
        heading_field = arcpy.Parameter(
            name="heading_field",
            displayName="Select Heading Field",
            datatype="Field",
            direction="Input",
            parameterType="Optional",
            category="Define Heading",
        )

        heading_field.parameterDependencies = [table_param.name]

        # Return parameter definitions as a list
        return [
            style_files,
            layers,
            table_param,
            field_param_1,
            field_param_2,
            field_param_3,
            field_param_4,
            label_field_1,
            label_delimiter,
            label_field_2,
            heading_field,
        ]

    def isLicensed(self):
        return True

    def updateParameters(self, parameters):
        return

    def updateMessages(self, _):
        return

    def execute(self, parameters, messages):
        # Retrieve a reference to the current project
        project = arcpy.mp.ArcGISProject("CURRENT")

        # Retrieve input parameters
        style_files = parameters[0].valueAsText
        input_layers = parameters[1].valueAsText
        table = parameters[2].valueAsText
        primary_key_field = parameters[3].valueAsText
        fill_symbol_code_field = parameters[4].valueAsText
        line_symbol_code_field = parameters[5].valueAsText
        point_symbol_code_field = parameters[6].valueAsText
        label_field_1 = parameters[7].valueAsText
        label_delimieter = parameters[8].valueAsText
        label_field_2 = parameters[9].valueAsText
        heading_field = parameters[10].valueAsText

        # Retrieve currently active map
        active_map = project.activeMap

        # Retrieve project layers
        map_layers = []
        if active_map:
            map_layers = active_map.listLayers()
        else:
            raise ValueError("Execution aborted: please select a map")

        # Check if selected layers are present in the active map
        input_layers_list = input_layers.split(";")
        layers_to_render = []
        for input_layer in input_layers_list:
            present = False
            for layer in map_layers:
                if layer.name == input_layer:
                    layers_to_render.append(layer.name)
                    present = True
                    break

            if not present:
                added_layer = active_map.addDataFromPath(input_layer)
                layers_to_render.append(added_layer.name)

        # Retrieve styles that are currently in the project
        project_styles = project.styles

        # Check if styles are in the project - add them if not present
        style_files_list = style_files.split(";")
        for style_path in style_files_list:
            if not style_path in project_styles:
                # Add the style to the project
                project_styles.append(style_path)
                project.updateStyles(project_styles)

        # Retrieve symbol codes from legend table
        symbol_codes = {}
        with arcpy.da.SearchCursor(
            table,
            [
                primary_key_field,
                fill_symbol_code_field,
                line_symbol_code_field,
                point_symbol_code_field,
            ],
        ) as search_cursor:
            for row in search_cursor:
                symbol_codes[row[0]] = {
                    "fill": row[1],
                    "stroke": row[2],
                    "marker": row[3],
                }

        # Retrieve label texts for custom labels
        add_labels = False
        if label_field_1 or label_field_2:
            add_labels = True
            custom_labels = {}
            labels = [primary_key_field]
            if label_field_1:
                labels.append(label_field_1)
            if label_field_2:
                labels.append(label_field_2)

            with arcpy.da.SearchCursor(
                table,
                labels,
            ) as search_cursor:
                for row in search_cursor:
                    custom_labels[row[0]] = f"{row[1]}{label_delimieter}{row[2]}"

        if heading_field:
            labels = [primary_key_field, heading_field]
            headings = {}
            with arcpy.da.SearchCursor(
                table,
                labels,
            ) as search_cursor:
                for row in search_cursor:
                    headings[row[0]] = row[1]

        # Start rendering process
        for layer in map_layers:
            if layer.name in layers_to_render:
                messages.AddMessage(f"Rendering layer: {layer.name}")

                sym = layer.symbology
                layer_desc = arcpy.Describe(layer)
                layer_shape_type = layer_desc.featureClass.shapeType

                # Check if UniqueValueRenderer is defined
                if sym.renderer == "UniqueValueRenderer":

                    # Regroup items according to heading field
                    if heading_field:
                        new_item_groups = {}
                        old_item_groups = {}

                        for group in sym.renderer.groups:
                            old_item_groups[group.heading] = group.items

                        for group in sym.renderer.groups:
                            for item in group.items:
                                if item.values[0][0] != "<Null>":
                                    leg_id = int(item.values[0][0])
                                else:
                                    leg_id = None

                                if leg_id in headings:
                                    heading = headings[leg_id]
                                    if not heading:
                                        heading = "Other"
                                else:
                                    heading = "Other"

                                if heading in new_item_groups:
                                    new_item_groups[heading].append(item)
                                else:
                                    new_item_groups[heading] = [item]

                        # Remove old item groups
                        sym.renderer.removeValues(old_item_groups)
                        layer.symbology = sym

                        # Add new item groups
                        sym.renderer.addValues(new_item_groups)
                        layer.symbology = sym

                    # Apply symbols from gallery
                    for group in sym.renderer.groups:
                        for item in group.items:
                            if item.values[0][0] != "<Null>":
                                leg_id = int(item.values[0][0])
                            else:
                                leg_id = None

                            if leg_id in symbol_codes:
                                symbol_code = get_symbol_code_for_shape(
                                    layer_shape_type, symbol_codes[leg_id]
                                )

                                if symbol_code == "#":
                                    continue

                                code_components = get_code_components(symbol_code)
                                symbol_key = code_components["symbol_key"]

                                if leg_id == 11804:
                                    messages.AddMessage(symbol_key)

                                if symbol_key:
                                    symbols_from_gallery = (
                                        item.symbol.listSymbolsFromGallery(symbol_key)
                                    )
                                    for symbol_from_gallery in symbols_from_gallery:
                                        if symbol_from_gallery.name == symbol_key:
                                            item.symbol = symbol_from_gallery
                                            break

                            # Add user defined labels
                            if add_labels:
                                if leg_id in custom_labels:
                                    item.label = custom_labels[leg_id]

                    layer.symbology = sym

                    # Retrieve CIM to add colors
                    cim_lyr = layer.getDefinition("V3")

                    for group in cim_lyr.renderer.groups:
                        for unique_value_class in group.classes:

                            if unique_value_class.values[0].fieldValues[0] != "<Null>":
                                leg_id = int(
                                    unique_value_class.values[0].fieldValues[0]
                                )
                            else:
                                leg_id = None

                            if leg_id in symbol_codes:
                                symbol_code = get_symbol_code_for_shape(
                                    layer_shape_type, symbol_codes[leg_id]
                                )
                                if symbol_code == "#":
                                    continue

                                code_components = get_code_components(symbol_code)
                                color_value = code_components["color"]
                                symbol_color_value = code_components["symbol_color"]

                                has_color = False
                                for (
                                    symbol_layer
                                ) in unique_value_class.symbol.symbol.symbolLayers:

                                    if symbol_color_value:
                                        update_symbol_layer_colors(
                                            symbol_layer, symbol_color_value
                                        )

                                    if color_value:
                                        # Update colors
                                        if layer_shape_type == "Polygon":
                                            if isinstance(
                                                symbol_layer,
                                                arcpy.cim.CIMSymbols.CIMSolidFill,
                                            ):
                                                update_color(symbol_layer, color_value)
                                                has_color = True

                                        elif layer_shape_type == "Polyline":
                                            if isinstance(
                                                symbol_layer,
                                                arcpy.cim.CIMSymbols.CIMSolidStroke,
                                            ):
                                                update_color(symbol_layer, color_value)
                                                has_color = True
                                        elif layer_shape_type == "Point":
                                            if isinstance(
                                                symbol_layer,
                                                arcpy.cim.CIMSymbols.CIMMarker,
                                            ):
                                                update_color(symbol_layer, color_value)
                                                has_color = True

                                if not has_color and color_value:
                                    if layer_shape_type == "Polygon":
                                        # Create CIMSolidFill
                                        fill_symbol = arcpy.cim.CIMSolidFill()
                                        update_color(fill_symbol, color_value)
                                        unique_value_class.symbol.symbol.symbolLayers.append(
                                            fill_symbol
                                        )
                                    elif layer_shape_type == "Polyline":
                                        # Create CIMSolidStroke
                                        stroke_symbol = arcpy.cim.CIMSolidStroke()
                                        update_color(stroke_symbol, color_value)
                                        unique_value_class.symbol.symbol.symbolLayers.append(
                                            stroke_symbol
                                        )
                                    elif layer_shape_type == "Point":
                                        # Create CIMMarker
                                        marker_symbol = arcpy.cim.CIMMarker()
                                        update_color(marker_symbol, color_value)
                                        unique_value_class.symbol.symbol.symbolLayers.append(
                                            marker_symbol
                                        )

                    # Push changes back to layer object
                    layer.setDefinition(cim_lyr)

                else:
                    raise ValueError(
                        f"Execution aborted: please apply a Unique Value Renderer to layer {layer.name}"
                    )

    def postExecute(self, _):
        return


# Update color property of symbol layer
def update_color(symbol_layer, color_value):
    color = arcpy.cim.CreateCIMObjectFromClassName("CIMCMYKColor", "V3")
    color.values = color_value["CMYK"]
    symbol_layer.color = color


# Update colors of symbol layers of type CIMCharacterMarker and CIMVectorMarker
def update_symbol_layer_colors(symbol_layer, symbol_color_value):
    if isinstance(
        symbol_layer,
        arcpy.cim.CIMSymbols.CIMCharacterMarker,
    ):
        for sub_symbol_layer in symbol_layer.symbol.symbolLayers:
            update_color(sub_symbol_layer, symbol_color_value)

    if isinstance(
        symbol_layer,
        arcpy.cim.CIMSymbols.CIMVectorMarker,
    ):
        for marker_graphic in symbol_layer.markerGraphics:
            for sub_symbol_layer in marker_graphic.symbol.symbolLayers:
                update_color(sub_symbol_layer, symbol_color_value)


# Retrieve symbole code for shape type
def get_symbol_code_for_shape(shape_type, symbol_codes):
    if shape_type == "Polygon":
        return symbol_codes["fill"]
    elif shape_type == "Polyline":
        return symbol_codes["stroke"]
    elif shape_type == "Point":
        return symbol_codes["marker"]
    else:
        raise ValueError(
            "Execute error: unknown shape type - allowed types: Point, Polyline and Polygon"
        )


# Decode symbol code
def get_code_components(code):
    code_len = len(code)
    if code_len == 4:
        return {
            "color": decode_color(code),
            "symbol_key": "",
            "symbol_color": {},
        }
    elif code_len == 10:
        return {
            "color": {},
            "symbol_key": code[0:7],
            "symbol_color": get_symbol_color(code[7:]),
        }
    elif code_len == 11:
        return {
            "color": decode_color(code[:4]),
            "symbol_key": code[4:12],
            "symbol_color": {},
        }
    elif code_len == 14:
        return {
            "color": decode_color(code[:4]),
            "symbol_key": code[4:11],
            "symbol_color": get_symbol_color(code[11:]),
        }
    else:
        raise ValueError(f"Execution aborted: unknown symbol code {code}")


# Decode color values
def decode_color(color_string):
    return {
        "CMYK": [
            get_percentage_from_letter(color_string[0]),
            get_percentage_from_letter(color_string[1]),
            get_percentage_from_letter(color_string[2]),
            get_percentage_from_letter(color_string[3]),
            100,
        ],
    }


# Decode letters
def get_percentage_from_letter(letter):
    if letter == "F":
        return 15
    elif letter == "V":
        return 100
    elif letter == "S":
        return 6
    elif letter == "X":
        return 0
    else:
        return int(letter) * 10


# Decode symbol colors
def get_symbol_color(color_string):
    if color_string == "BLK":
        return {
            "CMYK": [0, 0, 0, 100, 100],
        }
    elif color_string == "ROT":
        return {
            "CMYK": [0, 100, 100, 21, 100],
        }
    elif color_string == "CYN":
        return {
            "CMYK": [100, 0, 0, 0, 100],
        }
    elif color_string == "MGT":
        return {
            "CMYK": [0, 100, 0, 0, 100],
        }
    elif color_string == "BLA":
        return {
            "CMYK": [70, 30, 0, 21, 100],
        }
    elif color_string == "BRN":
        return {
            "CMYK": [45, 73, 93, 21, 100],
        }
    elif color_string == "GRN":
        return {
            "CMYK": [70, 0, 70, 21, 100],
        }
    elif color_string == "GLB":
        return {
            "CMYK": [0, 0, 100, 0, 100],
        }
    elif color_string == "ORA":
        return {
            "CMYK": [0, 32.16, 100, 0, 100],
        }
    elif color_string == "WHT":
        return {
            "CMYK": [100, 100, 100, 0, 100],
        }
    else:
        raise ValueError(f"Execution aborted: unknown color code {color_string}")