from __future__ import annotations

from collections import defaultdict
from typing import Any, Callable

from openpyxl.styles import Alignment, Border, Color, Font, PatternFill, Side
from openpyxl.worksheet.worksheet import Worksheet
from openpyxl.utils import get_column_letter

from src.services.export_service.export_mappings import PALETTE



RowNormalizer = Callable[[dict[str, Any]], dict[str, Any]]
ValueSerializer = Callable[[Any], Any]


class TabularSheetWriter:
    def __init__(
        self,
        columns: list[tuple[str, str]],
        row_normalizer: RowNormalizer,
        value_serializer: ValueSerializer,
    ):
        self.columns = columns
        self.row_normalizer = row_normalizer
        self.value_serializer = value_serializer

    def write(self, worksheet: Worksheet, rows: list[tuple]) -> None:
        worksheet.append([title for _, title in self.columns])
        for row in rows:
            raw_data = row[3] if len(row) > 3 and isinstance(row[3], dict) else {}
            normalized = self.row_normalizer(raw_data)
            worksheet.append(
                [self.value_serializer(normalized.get(key)) for key, _ in self.columns]
            )


class DepthSheetWriter:
    alignment = Alignment(horizontal='center', vertical='center')
    thin_border = Border(
        left=Side(style='thin'),
        right=Side(style='thin'),
        top=Side(style='thin'),
        bottom=Side(style='thin')
    )

    def __init__(
            self,
            columns_dict: dict,
            row_normalizer: RowNormalizer,
            value_serializer: ValueSerializer,
            colors_config: dict | None = None,
    ):
        self.columns_dict = columns_dict
        self._flat_columns_to_depth = {}
        self.columns = self._columns_dict_to_flat(columns_dict)
        self.row_normalizer = row_normalizer
        self.value_serializer = value_serializer
        self.colors_config = colors_config
        self.palette = None
        if colors_config:
            self._set_palette()
    
    def _columns_dict_to_flat(
            self, 
            columns_dict: dict | None = None, 
            prefix: str = "",
            depth_prefix: str = "",
    ) -> None:
        if columns_dict is None:
            columns_dict = self.columns_dict
        result = []
        for key, sub_dict in columns_dict.items():
            if "key" in sub_dict:
                new_key = sub_dict['key']
            else:
                new_key = key
            if "children" not in sub_dict:
                if new_key:
                    add_key = f"{prefix}{new_key}"
                    result.append(add_key)
                    self._flat_columns_to_depth[add_key] = f"{depth_prefix}{key}"
            else:
                result += self._columns_dict_to_flat(
                    sub_dict["children"], 
                    prefix=f"{prefix}{new_key}." if new_key else prefix,
                    depth_prefix=f"{depth_prefix}{key}.",
                )
        return result
    
    def _set_palette(self) -> None:
        self.palette = {}
        for key in self.columns:
            full_key = None
            for sub_key in self._flat_columns_to_depth[key].split("."):
                if full_key is None:
                    full_key = sub_key
                else:
                    full_key = f"{full_key}.{sub_key}"
                if full_key in self.colors_config and "palette" in self.colors_config[full_key]:
                    self.palette[key] = self.colors_config[full_key]["palette"]

    def _get_header_max_row(self, column_dict: dict) -> int:
        max_depth = 0
        for sub_dict in column_dict.values():
            if "children" in sub_dict:
                max_depth = max(max_depth, self._get_header_max_row(sub_dict["children"]))
        return max_depth + 1

    def _write_header(
            self, 
            worksheet: Worksheet, 
            columns_dict: dict | None = None,
            current_column: int = 1,
            current_row: int = 1,
            max_row: int | None = None,
            full_key: str | None = None
    ) -> int:
        if columns_dict is None:
            columns_dict = self.columns_dict
            max_row = self._get_header_max_row(self.columns_dict)
        for key, sub_dict in columns_dict.items():
            if full_key:
                key = f"{full_key}.{key}"

            if "children" not in sub_dict:
                cell = worksheet.cell(row=current_row, column=current_column, value=sub_dict["name"])
                cell.border = self.thin_border
                if max_row != current_row:
                    worksheet.merge_cells(
                        start_row=current_row, 
                        start_column=current_column, 
                        end_row=max_row, 
                        end_column=current_column,
                    )
                current_column += 1
            else:
                new_column = self._write_header(
                    worksheet=worksheet,
                    columns_dict=sub_dict["children"],
                    current_column=current_column,
                    current_row=current_row+1,
                    max_row=max_row,
                    full_key=key,
                )
                cell = worksheet.cell(row=current_row, column=current_column, value=sub_dict["name"])
                if new_column - current_column > 1:
                    worksheet.merge_cells(
                        start_row=current_row, 
                        start_column=current_column, 
                        end_row=current_row,
                        end_column=new_column - 1,
                    )
                cell.border = self.thin_border
                current_column = new_column
            cell.alignment = self.alignment
            
            if self.colors_config and (color := self.colors_config.get(key)):
                if "background" in color:
                    cell.fill = PatternFill(
                        start_color=Color(rgb=color["background"]),
                        fill_type="solid",
                    )
                if "color" in color:
                    cell.font = Font(color=color["color"])
        return current_column

    def _set_widths(self, worksheet: Worksheet):
        for col in worksheet.columns:
            max_len = 0
            col_letter = get_column_letter(col[0].column)  # Get alphabetical column name (e.g., 'A')
            
            for cell in col:
                if cell.value is not None:
                    cell_len = len(str(cell.value))
                    if cell_len > max_len:
                        max_len = cell_len
                        
            adjusted_width = max(max_len + 3, 10)
            worksheet.column_dimensions[col_letter].width = adjusted_width

    def _set_color_to_row(self, ws: Worksheet, item_type: str) -> None:
        if not self.palette or item_type in ("SUB_ITEM", "INPUT"):
            return

        row = ws.max_row

        for i, key in enumerate(self.columns):
            if key not in self.palette:
                continue
            cell = ws.cell(column=i+1, row=row)
            if (palette := self.palette.get(key)):
                cell.fill = PatternFill(
                    start_color=Color(rgb=PALETTE[palette][item_type]),
                    fill_type="solid",
                )

    def write(self, worksheet: Worksheet, rows: list[tuple]) -> None:
        self._write_header(worksheet)
        for row in rows:
            raw_data = row[3] if len(row) > 3 and isinstance(row[3], dict) else {}
            normalized = self.row_normalizer(raw_data)
            worksheet.append(
                [self.value_serializer(normalized.get(key)) for key in self.columns]
            )
            if len(row) > 3:
                self._set_color_to_row(ws=worksheet, item_type=row[0])
        self._set_widths(worksheet)
