from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Sequence

from openpyxl.worksheet.worksheet import Worksheet

from xllayout.core.canvas import Anchor, Box
from xllayout.core.cell_style import CellStyleSpec, apply_cell_style, merge_cell_styles
from xllayout.core.theme import Theme
from xllayout.core.utils import set_col_widths


TableStyle = str


@dataclass
class TableBlock:
    headers: list[str]
    rows: list[Sequence[Any]]
    table_style: TableStyle = "zebra_compact"  # simple | zebra_compact | kpi_table
    span_cols: int | None = None
    height_rows: int | None = None
    col_widths: list[float] | None = None
    number_formats_by_col: dict[int, str] = field(default_factory=dict)  # 1-based column idx
    show_header: bool = True

    # New style API (v1.1.0)
    header_style: CellStyleSpec | None = None
    body_style: CellStyleSpec | None = None
    header_style_by_col: dict[int, CellStyleSpec] = field(default_factory=dict)
    body_style_by_col: dict[int, CellStyleSpec] = field(default_factory=dict)
    wrap_text_cols: set[int] = field(default_factory=set)
    header_wrap_text_cols: set[int] = field(default_factory=set)
    body_wrap_text_cols: set[int] = field(default_factory=set)

    # Optional row heights
    header_row_height: float | None = None
    default_row_height: float | None = None
    row_height_by_row: dict[int, float] = field(default_factory=dict)  # 1-based body row index

    def measure(self, anchor: Anchor, span_rows: int | None = None, span_cols: int | None = None) -> Box:
        data_col_count = max([len(self.headers)] + [len(row) for row in self.rows] + [1])
        requested_width = span_cols if span_cols is not None else self.span_cols
        width = max(data_col_count, requested_width or data_col_count)

        data_height = (1 if self.show_header else 0) + len(self.rows)
        requested_height = span_rows if span_rows is not None else self.height_rows
        height = max(1, requested_height or data_height)

        return Box(anchor.row, anchor.col, anchor.row + height - 1, anchor.col + width - 1)

    def draw(self, ws: Worksheet, box: Box, theme: Theme) -> None:
        if self.col_widths:
            set_col_widths(ws, box.c1, self.col_widths[: box.width])

        cursor_row = box.r1
        if self.show_header and cursor_row <= box.r2:
            self._draw_header(ws, box, theme, cursor_row)
            if self.header_row_height is not None:
                ws.row_dimensions[cursor_row].height = self.header_row_height
            cursor_row += 1

        max_data_rows = box.r2 - cursor_row + 1
        if max_data_rows < 0:
            max_data_rows = 0

        for idx, row_values in enumerate(self.rows[:max_data_rows]):
            row_no = cursor_row + idx
            self._draw_row(ws, box, theme, row_no, row_values, idx)

            row_height = self.row_height_by_row.get(idx + 1, self.default_row_height)
            if row_height is not None:
                ws.row_dimensions[row_no].height = row_height

        if not self.rows and cursor_row <= box.r2:
            cell = ws.cell(row=cursor_row, column=box.c1, value="No rows available for this section")
            cell.font = theme.fonts.small
            cell.alignment = theme.alignments.left
            cell.border = theme.borders.thin

    def _draw_header(self, ws: Worksheet, box: Box, theme: Theme, row: int) -> None:
        for i in range(box.width):
            col_idx = i + 1
            col = box.c1 + i
            value = self.headers[i] if i < len(self.headers) else None
            cell = ws.cell(row=row, column=col, value=value)
            cell.font = theme.fonts.section
            cell.fill = theme.fills.header
            cell.border = theme.borders.thin
            cell.alignment = theme.alignments.wrap_center

            style = self._header_style_for_col(col_idx)
            apply_cell_style(
                cell,
                style,
                default_font=theme.fonts.section,
                default_fill=theme.fills.header,
                default_border=theme.borders.thin,
                default_alignment=theme.alignments.wrap_center,
            )

    def _draw_row(
        self,
        ws: Worksheet,
        box: Box,
        theme: Theme,
        row_no: int,
        row_values: Sequence[Any],
        row_index: int,
    ) -> None:
        zebra = self.table_style == "zebra_compact" and row_index % 2 == 1
        for i in range(box.width):
            col_idx = i + 1
            col = box.c1 + i
            value = row_values[i] if i < len(row_values) else None
            cell = ws.cell(row=row_no, column=col, value=value)

            if self.table_style == "kpi_table" and i == 0:
                cell.font = theme.fonts.section
            else:
                cell.font = theme.fonts.body

            if zebra:
                cell.fill = theme.fills.zebra

            cell.border = theme.borders.thin
            cell.alignment = self._pick_alignment(theme, col_idx, value)

            fmt = self.number_formats_by_col.get(col_idx)
            if fmt:
                cell.number_format = fmt
            elif isinstance(value, int):
                cell.number_format = theme.number_formats.integer

            style = self._body_style_for_col(col_idx)
            apply_cell_style(
                cell,
                style,
                default_border=theme.borders.thin,
                default_alignment=self._pick_alignment(theme, col_idx, value),
            )

    def _pick_alignment(self, theme: Theme, col_idx: int, value: Any):
        if col_idx in self.number_formats_by_col:
            if self.number_formats_by_col[col_idx] == theme.number_formats.percent:
                return theme.alignments.right
            return theme.alignments.right
        if isinstance(value, (int, float)):
            return theme.alignments.right
        return theme.alignments.left

    def _header_style_for_col(self, col_idx: int) -> CellStyleSpec | None:
        wrap_style = self._wrap_style(enabled=(col_idx in self.wrap_text_cols or col_idx in self.header_wrap_text_cols))
        return merge_cell_styles(self.header_style, self.header_style_by_col.get(col_idx), wrap_style)

    def _body_style_for_col(self, col_idx: int) -> CellStyleSpec | None:
        wrap_style = self._wrap_style(enabled=(col_idx in self.wrap_text_cols or col_idx in self.body_wrap_text_cols))
        return merge_cell_styles(self.body_style, self.body_style_by_col.get(col_idx), wrap_style)

    @staticmethod
    def _wrap_style(enabled: bool) -> CellStyleSpec | None:
        if not enabled:
            return None
        return CellStyleSpec(wrap_text=True)
