"""`workflow` 共享输出资源: `workbook` 实现(内部模块).

说明:
- 承载 `workbook_sheet`/`workbook_append` 的计划构建、对齐与提交落盘
- 运行时需兼容 `Python 3.6`
"""

from abc import ABC
from contextlib import suppress
from pathlib import Path
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, cast

from .._internal.utils.excel import escape_excel_formula
from ..events import EVENT_DIAGNOSTIC_WARNING
from ..events._events import DiagnosticWarningEvent
from ..sinks._internal.base import create_temp_path
from ..vendor.compact.importlibx import require_optional_dependency
from ..vendor.compact.typing_extensionsx import override
from ..vendor.dataclassesx import dataclass
from .resources_base import ScalimWorkflowWriteError, WorkflowResourceManagerBase
from .resources_csv import AppendSegment, WorkflowCsvInput, build_alignment_mapping, describe_header_diff, iter_csv_rows, read_csv_header

# 内部实现仍沿用原有局部命名,减少重构噪音.
_AppendSegment = AppendSegment
_build_alignment_mapping = build_alignment_mapping
_describe_header_diff = describe_header_diff
_iter_csv_rows = iter_csv_rows
_read_csv_header = read_csv_header

if TYPE_CHECKING:
    from openpyxl import Workbook


def _get_openpyxl_workbook_class() -> "Type[Workbook]":
    openpyxl_mod = require_optional_dependency("openpyxl", context="scalim.workflow.resources")
    return cast("Any", openpyxl_mod).Workbook  # pragma: allow-cast openpyxl module runtime boundary


def _best_effort_close_write_only_workbook_worksheets(workbook: Any) -> None:
    try:
        worksheets_obj = workbook.worksheets
    except AttributeError:
        return
    try:
        worksheets = list(worksheets_obj)
    except TypeError:
        return
    for ws in worksheets:
        try:
            is_closed = bool(ws.closed)
        except AttributeError:
            continue
        if is_closed:
            continue
        with suppress(Exception):
            ws.close()


@dataclass
class _SheetPlan:
    sheet: str
    baseline_header: List[str]
    segments: List[_AppendSegment]


@dataclass
class _WorkbookPlan:
    resource_id: str
    path: str
    allow_formulas: bool
    sheet_decl_order: Dict[str, int]
    sheet_order: List[str]
    sheets: Dict[str, _SheetPlan]
    last_workflow_node_id: Optional[str] = None


class _WorkflowWorkbookResourceMixin(WorkflowResourceManagerBase, ABC):
    def _get_or_create_workbook(self, workbook_id: str, *, workflow_node_id: str) -> _WorkbookPlan:
        key = str(workbook_id)

        def _create() -> object:
            raw_path = self._workbook_defs.get(key)
            if raw_path is None:
                msg = "Unknown workbook resource id: {!r}".format(key)
                raise ScalimWorkflowWriteError(msg)
            return _WorkbookPlan(
                resource_id=key,
                path=str(raw_path),
                allow_formulas=bool(self._workbook_allow_formulas.get(key, False)),
                sheet_decl_order={},
                sheet_order=[],
                sheets={},
            )

        def _on_create(plan: object) -> None:
            p = cast("_WorkbookPlan", plan)  # pragma: allow-cast joinable plan typed narrowing
            self._emit_resource_create(
                workflow_node_id=str(workflow_node_id),
                resource_type="workbook",
                resource_id=key,
                path=str(p.path),
            )

        plan = self._get_or_create_joinable_plan(
            resource_type="workbook",
            resource_id=key,
            plans=self._workbooks,
            inflight=self._inflight_workbooks,
            create_fn=_create,
            on_create=_on_create,
        )
        return cast("_WorkbookPlan", plan)  # pragma: allow-cast joinable plan typed narrowing

    def apply_workbook_sheet(
        self,
        *,
        workflow_node_id: str,
        decl_order: int,
        workbook_id: str,
        sheet: str,
        input_node_id: str,
        input_output_id: str,
        input_csv: WorkflowCsvInput,
        on_conflict: str,
    ) -> None:
        plan = self._get_or_create_workbook(workbook_id, workflow_node_id=str(workflow_node_id))
        sheet_name = str(sheet)
        action = "write"

        input_header = _read_csv_header(input_csv)

        pending_skip = False
        with self._lock:
            existing = plan.sheets.get(sheet_name)
            if existing is not None:
                if on_conflict == "skip":
                    action = "skip"
                    plan.last_workflow_node_id = str(workflow_node_id)
                    pending_skip = True
                if on_conflict == "error":
                    msg = "Sheet conflict (workbook_sheet): workbook={!r}, sheet={!r}".format(str(workbook_id), sheet_name)
                    raise ScalimWorkflowWriteError(msg, diff=["on_conflict=error", "existing_sheet=present"])
                if on_conflict == "overwrite":
                    action = "overwrite"

            if not pending_skip:
                existing_decl_order = plan.sheet_decl_order.get(sheet_name)
                resolved_decl_order = int(decl_order)
                if existing_decl_order is None or resolved_decl_order < int(existing_decl_order):
                    plan.sheet_decl_order[sheet_name] = resolved_decl_order
                plan.sheet_order = sorted(plan.sheet_decl_order.keys(), key=lambda name: (plan.sheet_decl_order.get(name, 0), str(name)))

                mapping = _build_alignment_mapping(input_header, input_header)
                segment = _AppendSegment(
                    decl_order=int(decl_order),
                    input_csv=input_csv,
                    header_policy="once",
                    mapping=mapping,
                    on_mismatch="error",
                    align_by="header",
                    input_header=input_header,
                )
                plan.sheets[sheet_name] = _SheetPlan(sheet=sheet_name, baseline_header=list(input_header), segments=[segment])
                plan.last_workflow_node_id = str(workflow_node_id)

        if pending_skip:
            self._emit_resource_write(
                workflow_node_id=str(workflow_node_id),
                resource_type="workbook",
                resource_id=str(workbook_id),
                path=str(plan.path),
                write_kind="workbook_sheet",
                action=action,
                input_node_id=str(input_node_id),
                input_output_id=str(input_output_id),
                sheet=sheet_name,
            )
            return

        self._emit_resource_write(
            workflow_node_id=str(workflow_node_id),
            resource_type="workbook",
            resource_id=str(workbook_id),
            path=str(plan.path),
            write_kind="workbook_sheet",
            action=action,
            input_node_id=str(input_node_id),
            input_output_id=str(input_output_id),
            sheet=sheet_name,
        )

    def apply_workbook_append(
        self,
        *,
        workflow_node_id: str,
        decl_order: int,
        workbook_id: str,
        sheet: str,
        input_node_id: str,
        input_output_id: str,
        input_csv: WorkflowCsvInput,
        align_by: str,
        header_policy: str,
        on_mismatch: str,
    ) -> None:
        plan = self._get_or_create_workbook(workbook_id, workflow_node_id=str(workflow_node_id))
        sheet_name = str(sheet)
        input_header = _read_csv_header(input_csv)

        pending_warning: Optional[DiagnosticWarningEvent] = None
        pending_warning_meta: Optional[Dict[str, object]] = None
        pending_skip = False

        with self._lock:
            sheet_plan = plan.sheets.get(sheet_name)
            if sheet_plan is None:
                plan.sheet_decl_order[sheet_name] = int(decl_order)
                plan.sheet_order = sorted(plan.sheet_decl_order.keys(), key=lambda name: (plan.sheet_decl_order.get(name, 0), str(name)))
                sheet_plan = _SheetPlan(sheet=sheet_name, baseline_header=list(input_header), segments=[])
                plan.sheets[sheet_name] = sheet_plan
            else:
                existing_decl_order = plan.sheet_decl_order.get(sheet_name)
                resolved_decl_order = int(decl_order)
                if existing_decl_order is None or resolved_decl_order < int(existing_decl_order):
                    plan.sheet_decl_order[sheet_name] = resolved_decl_order
                    plan.sheet_order = sorted(
                        plan.sheet_decl_order.keys(),
                        key=lambda name: (plan.sheet_decl_order.get(name, 0), str(name)),
                    )

            expected = list(sheet_plan.baseline_header)
            mapping = _build_alignment_mapping(expected, input_header)

            if list(input_header) != expected:
                diff = _describe_header_diff(expected, input_header)
                if on_mismatch == "error":
                    msg = "Field alignment mismatch (workbook_append): workbook={!r}, sheet={!r}".format(str(workbook_id), sheet_name)
                    raise ScalimWorkflowWriteError(msg, diff=diff)
                if on_mismatch == "warn":
                    pending_warning = DiagnosticWarningEvent(
                        message="Field alignment mismatch (warn): workbook={!r}, sheet={!r}".format(str(workbook_id), sheet_name),
                        source_id=None,
                        field_id=None,
                        lookup_key={"expected": expected, "actual": list(input_header)},
                        row_id=None,
                    )
                    pending_warning_meta = {"workflow_exec_id": self._workflow_exec_id, "workflow_node_id": str(workflow_node_id)}
                if on_mismatch == "skip":
                    plan.last_workflow_node_id = str(workflow_node_id)
                    pending_skip = True

            if not pending_skip:
                sheet_plan.segments.append(
                    _AppendSegment(
                        decl_order=int(decl_order),
                        input_csv=input_csv,
                        header_policy=str(header_policy),
                        mapping=mapping,
                        on_mismatch=str(on_mismatch),
                        align_by=str(align_by),
                        input_header=list(input_header),
                    )
                )
                plan.last_workflow_node_id = str(workflow_node_id)

        if pending_warning is not None:
            _ = self._instrumentation.emit(EVENT_DIAGNOSTIC_WARNING, pending_warning, meta=pending_warning_meta)

        if pending_skip:
            self._emit_resource_write(
                workflow_node_id=str(workflow_node_id),
                resource_type="workbook",
                resource_id=str(workbook_id),
                path=str(plan.path),
                write_kind="workbook_append",
                action="skip",
                input_node_id=str(input_node_id),
                input_output_id=str(input_output_id),
                sheet=sheet_name,
            )
            return

        self._emit_resource_write(
            workflow_node_id=str(workflow_node_id),
            resource_type="workbook",
            resource_id=str(workbook_id),
            path=str(plan.path),
            write_kind="workbook_append",
            action="append",
            input_node_id=str(input_node_id),
            input_output_id=str(input_output_id),
            sheet=sheet_name,
        )

    @override
    def _commit_workbook(self, plan: object) -> None:  # noqa: C901
        p = cast("_WorkbookPlan", plan)  # pragma: allow-cast joinable plan typed narrowing
        if not p.sheets:
            return

        final_path = str(p.path)
        staging_path = self._staging_path_for_final_output(final_path)

        try:
            workbook_cls = _get_openpyxl_workbook_class()
        except ImportError as exc:
            raise ScalimWorkflowWriteError(str(exc)) from exc

        wb = workbook_cls(write_only=True)
        try:
            for sheet_name in p.sheet_order:
                sheet_plan = p.sheets.get(sheet_name)
                if sheet_plan is None:  # pragma: no cover  # pragma: allow-no-cover unreachable: sheet_plan always exists
                    continue  # pragma: no cover  # pragma: allow-no-cover unreachable: sheet_plan always exists
                ws = wb.create_sheet(str(sheet_name))
                header_written = False
                segments = sorted(sheet_plan.segments, key=lambda seg: int(seg.decl_order))
                for seg in segments:
                    if seg.header_policy == "always" or (seg.header_policy == "once" and not header_written):
                        header = [escape_excel_formula(x, allow_formulas=bool(p.allow_formulas)) for x in list(sheet_plan.baseline_header)]
                        _ = ws.append(header)
                        header_written = True
                    # `never`: 不输出 `header`

                    for row in _iter_csv_rows(seg.input_csv):
                        out_row: List[object] = []
                        for idx in seg.mapping:
                            value = row[idx] if idx >= 0 and idx < len(row) else ""
                            out_row.append(escape_excel_formula(value, allow_formulas=bool(p.allow_formulas)))
                        _ = ws.append(out_row)

            temp_path = create_temp_path(staging_path, ".xlsx.tmp")
            temp_obj = Path(temp_path)
            try:
                wb.save(temp_obj)
                _ = temp_obj.replace(staging_path)
            except Exception as exc:
                with suppress(Exception):
                    temp_obj.unlink()
                msg = "Workbook commit failed: {}: {}".format(type(exc).__name__, exc)
                raise ScalimWorkflowWriteError(msg) from exc
        except Exception:
            _best_effort_close_write_only_workbook_worksheets(wb)
            raise
        finally:
            with suppress(Exception):
                wb.close()

        node_id = p.last_workflow_node_id or "__wf__commit"
        self._register_staged_output(
            resource_type="workbook",
            resource_id=p.resource_id,
            workflow_node_id=str(node_id),
            staged_path=str(staging_path),
            final_path=str(final_path),
        )

    @override
    def _discard_workbook(self, plan: object, *, workflow_node_id: str, reason: str) -> None:
        p = cast("_WorkbookPlan", plan)  # pragma: allow-cast joinable plan typed narrowing
        node_id = p.last_workflow_node_id or str(workflow_node_id)
        self._emit_resource_discard(
            workflow_node_id=node_id,
            resource_type="workbook",
            resource_id=p.resource_id,
            path=str(p.path),
            reason=str(reason),
        )


__all__ = (
    "SheetPlan",
    "WorkbookPlan",
    "WorkflowWorkbookResourceMixin",
    "best_effort_close_write_only_workbook_worksheets",
    "get_openpyxl_workbook_class",
)

SheetPlan = _SheetPlan
WorkbookPlan = _WorkbookPlan
WorkflowWorkbookResourceMixin = _WorkflowWorkbookResourceMixin
best_effort_close_write_only_workbook_worksheets = _best_effort_close_write_only_workbook_worksheets
get_openpyxl_workbook_class = _get_openpyxl_workbook_class
