fledermaus/fledermaus_formatter/table_handler/serializer.py

from openpyxl import Workbook, load_workbook
from pprint import pprint as print

import fledermaus_formatter.instructions.mappings as mappings
from fledermaus_formatter.table_handler.Entry import Entry

from .commands.find_and_replace import FindAndReplace
from .commands.regex_edit_col import RegexEditCol
from .commands.regex_move_to_col import RegexMoveToCol
from .commands.copy import Copy
from .commands.regex_group_valuebased_choose import RegexGroupValuebasedChoose
from .commands.row_color_to_value import RowColorToCalue
from fledermaus_formatter.logger import logger

from .conditions.regex_cell import RegexCellCondition

def hex_to_rgb(hex):
    if type(hex) != str and str("Values must be of type <class 'str'>"):
        # chatch wierd openxl things
        hex = "#0"
    hex=  "{:06}".format(str(hex).lstrip("#"))
    return tuple(int(hex[i:i+2],16) for i in (0, 2, 4))

class TableSerializer:
    def __init__(self, file: str, expected_input_schema: list[str], expected_output_schema: list[str]):
        self.wb: Workbook = load_workbook(file, read_only=True)
        self.ws = self.wb.active
        if not self.ws:
            raise ValueError("No active worksheet found in the workbook.")
        
        self.file = file
        self.data = self._serialize_table(expected_input_schema)
        self.serialized_output = {}
        self.seperator = ", "
        self.expected_output_schema = expected_output_schema or expected_input_schema
        self.implicit_copy = True
        self.output_header = True


    def _serialize_table(self, expected_cols: list[str]) -> dict[int, list[Entry]]:
        data = {col: {} for col in expected_cols}
        map_expected_with_ws = {}
        for idx, cell in enumerate(self.ws[1]):
            cell_value = cell.value if cell.value else ""
            if cell_value in expected_cols:
                map_expected_with_ws[cell_value] = idx
            else:
                print(f"Skipping unexpected Col {cell_value}")

        if set(map_expected_with_ws.keys()) != set(expected_cols):
            raise ValueError("Missing Col in input, check your table header")

        for rid, row in enumerate(self.ws.iter_rows(min_row=2, values_only=False), start=2):
            for expected_col, ws_idx in map_expected_with_ws.items():
                cell = row[ws_idx]
                data[expected_col][rid] = Entry(
                    value=str(cell._value) if not cell.value is None else "",
                    bg_color=hex_to_rgb(cell.fill.start_color.rgb if cell.fill and cell.fill.start_color else "#0"),
                    fg_color=hex_to_rgb(cell.font.color.rgb if cell.font and cell.font.color else "#0")
                )
        return data
    
    def write_output(self, destination_file: str):
        if not self.serialized_output:
            raise ValueError("There is nothing to write, serialize first")
        
        # write header
        owb = Workbook()
        ows = owb.active

        if self.output_header:
            for col_i, col_name in enumerate(self.expected_output_schema, start=1):
                ows.cell(row=1, column=col_i, value=col_name)

        # write rows
        for row_i, (row_id, row_data) in enumerate(self.serialized_output.items(), start=2 if self.output_header else 1):
            for col_i, col_name in enumerate(self.expected_output_schema, start=1):
                ows.cell(row=row_i, column=col_i, value=row_data.get(col_name).value if row_data.get(col_name) else "")
        
        owb.save(destination_file)

    
    def _check_conditions(self, instruct: dict[str, str], row: dict[str, any]):
        when_regex_condition_map = instruct.get("when_regex")
        if when_regex_condition_map:
            return RegexCellCondition(when_regex_condition_map).check_on_row(row)
        return True

    def _merge_change_to_cel(self, value, changes):
        #if not value and changes:
        #    value = changes
        #elif changes:
        #    value = f"{value}{self.seperator}{changes}"
        value = changes
        return value

    def _extecute_instruct(self, source_cell: str, instruct: dict, row: dict[str, any], working_row: dict[str, any]):
        in_place_working_row = {**row, **working_row}
        
        # execute changes
        res = None
        res_row = {}
        match instruct.get("command"):
            case "FIND_AND_REPLACE":
                res = FindAndReplace(instruct["map"], instruct.get("reqire_full_match", False)).execute_on_cell(in_place_working_row[source_cell])
            case "REGEX_MOVE_TO_COL":
                res_row = RegexMoveToCol(instruct["match"], instruct["to"], instruct.get("with_input_iteration")) \
                    .execute_on_row(in_place_working_row, source_cell)
            case "REGEX_EDIT_COL":
                res = RegexEditCol(instruct["match"], instruct["replace_with"], instruct.get("matching_group_rearange"))\
                    .execute_on_cell(in_place_working_row[source_cell])
            case "REGEX_GROUP_VALUEBASED_CHOOS":
                res = RegexGroupValuebasedChoose(instruct["match"], instruct["highest"])\
                    .execute_on_cell(in_place_working_row[source_cell])
            case "CONDITIONAL_COPY":
                res_row = Copy(instruct["source"], instruct["target"], replace=instruct.get("replace"))\
                    .execute_on_row(in_place_working_row)
            case "ROW_COLOR_TO_BOOL":
                res_row = RowColorToCalue(instruct["target"], instruct["value"], instruct.get("fg_colors"), instruct.get("bg_colors"), instruct.get("approximate"), instruct.get("threshold"))\
                    .execute_on_row(in_place_working_row)

        # Merge row_changes to working_row
        for target_cel, change in res_row.items():
            working_row[target_cel] = self._merge_change_to_cel(working_row.get(target_cel), change)

        # Merge changes in working row
        if source_cell and res is not None:
            working_row[source_cell] = self._merge_change_to_cel(working_row.get(source_cell), res)
        return working_row


    @property
    def vertical_len(self):
        return max(len(data_row) for data_row in self.data.values())

    def get_row_as_dict(self, rid: int):
        return {col: rows.get(rid) for col, rows in self.data.items()}

    def run(self, col_instructions: dict[tuple[str], dict], row_instructions: dict[tuple[str], dict]):
        self.serialized_output = {}
        amount_working_rows = self.vertical_len +2
        for rid in range(2, amount_working_rows):
            if rid % 1000 == 0:
                logger.info(f"{self.file} Working on row {rid}/{amount_working_rows} ")
            working_row = {}
            source_row = self.get_row_as_dict(rid)
            for source_cols, instructions in list(col_instructions.items()) + [
                    ((None,), instr) for instr in row_instructions
                ]:
                for instruct in instructions if type(instructions) == list else [instructions]:
                    for col in source_cols:
                        if not self._check_conditions(instruct, source_row):
                            continue
                        working_row = self._extecute_instruct(col, instruct, source_row, working_row)
            if not self.implicit_copy:
                continue
            for col in self.expected_output_schema:
                if col in working_row:
                    continue
                working_row[col] = source_row.get(col)
            self.serialized_output[rid] = working_row