from pathlib import Path

from lxml.etree import Element, XMLParser, fromstring, tostring

PARSER = XMLParser(remove_blank_text=False)
BASE_DIR = Path("charts")


def read_svg(path: Path) -> Element:
    raw = path.read_bytes()

    return fromstring(raw, PARSER)


def write_svg(path: Path, element: Element) -> None:
    raw = tostring(element, encoding="utf-8", pretty_print=False, xml_declaration=True)

    path.write_bytes(raw)


def split_plotly_svg(input_path: Path, outline_path: Path, data_path: Path) -> None:
    outline_tree, data_tree = read_svg(input_path), read_svg(input_path)

    data_xpath = "//*[local-name()='g' and contains(@class, 'overplot')]"
    for element in outline_tree.xpath(data_xpath):
        element.attrib["display"] = "none"

    write_svg(outline_path, outline_tree)

    frame_layers = [
        "bglayer",
        "gridlayer",
        "zerolinelayer",
        "axis",
        "infolayer",
        "shapelayer",
        "overaxes-above",
    ]
    frame_conditions = " or ".join(
        [f"contains(@class, '{cls}')" for cls in frame_layers]
    )
    frame_xpath = f"//*[local-name()='g' and ({frame_conditions})]"
    for element in data_tree.xpath(frame_xpath):
        element.attrib["display"] = "none"

    bg_rect_xpath = "//*[local-name()='rect' and (contains(@class, 'bg') or contains(@class, 'plot-bg'))]"
    for rect in data_tree.xpath(bg_rect_xpath):
        rect.attrib["display"] = "none"

    write_svg(data_path, data_tree)


def split_all_svgs() -> None:
    paths = list(BASE_DIR.glob("*.svg"))

    for path in paths:
        if path.stem.endswith(("_outline", "_data")):
            continue

        split_plotly_svg(
            path,
            path.with_stem(f"{path.stem}_outline"),
            path.with_stem(f"{path.stem}_data"),
        )