from __future__ import annotations

from pathlib import Path
from typing import cast

from openpyxl import Workbook
from openpyxl.worksheet.worksheet import Worksheet

from ._parser import XmlSource, load_xml_root
from ._renderer import build_sheet


def generate_workbook_from_xml(xml_source: XmlSource) -> Workbook:
    """Generate an in-memory Workbook from an XML definition."""
    root = load_xml_root(xml_source)
    if root.tag != "workbook":
        raise ValueError("Root element must be <workbook>")

    workbook = Workbook()
    sheets = list(root.findall("sheet"))

    if not sheets:
        return workbook

    for index, sheet_node in enumerate(sheets, start=1):
        title = sheet_node.get("name") or f"Sheet{index}"
        if index == 1:
            ws = cast(Worksheet, workbook.active)
            ws.title = title
        else:
            ws = workbook.create_sheet(title=title)

        build_sheet(ws, sheet_node)

    return workbook


def write_workbook_from_xml(
    xml_source: XmlSource,
    *,
    save_to: str | Path,
) -> Workbook:
    """Write the generated workbook to disk and return it."""
    workbook = generate_workbook_from_xml(xml_source)
    workbook.save(str(save_to))
    return workbook
