#!/usr/bin/env python3
from pathlib import Path
import re

ROOT = Path(__file__).resolve().parents[1]
TEMPLATE_ROOT = ROOT / "modules/chart_engine/template/d3-js"


def subn(pattern: str, repl, text: str, count: int = 0) -> tuple[str, int]:
    return re.subn(pattern, repl, text, count=count, flags=re.MULTILINE)


def replace_capsule_rows(text: str) -> tuple[str, int]:
    if "legendCapsuleHeightH" not in text or "legendRows.forEach" not in text:
        return text, 0
    if "chartUtils.legend.draw" in text and "markerShape: \"capsule\"" in text:
        return text, 0
    pattern = (
        r"\n\s*const legendGroup = svg\.append\(\"g\"\)\s*\n"
        r"\s*\.attr\(\"class\", \"legend\"\)[\s\S]*?"
        r"\s*currentY \+= rowHeight;\s*\n"
        r"\s*\}\);\s*\n"
    )
    replacement = (
        "\n    chartUtils.legend.draw(svg, groups, {\n"
        "        color: (group, index) => `url(#pattern-${group.replace(/[^a-zA-Z0-9]/g, '-')}-${index})`,\n"
        "        colorResolver,\n"
        "        colors,\n"
        "    }, {\n"
        "        y: margin.top * 0.7,\n"
        "        maxWidth: width - 100,\n"
        "        align: \"center\",\n"
        "        markerShape: \"capsule\",\n"
        "        markerWidth: legendCapsuleWidthH || 15,\n"
        "        markerHeight: legendCapsuleHeightH || 10,\n"
        "        labelGap: legendTextPadding || 5,\n"
        "        itemGap: legendItemPadding || 15,\n"
        "        rowGap: legendVerticalSpacing || 10,\n"
        "        fontSize: typography.label.font_size,\n"
        "        fontFamily: typography.label.font_family,\n"
        "        fontWeight: typography.label.font_weight,\n"
        "        textColor: colorResolver.text({ fallback: \"#333333\" }).value,\n"
        "    });\n"
    )
    text2, n = subn(pattern, replacement, text, count=1)
    if n:
        text2 = re.sub(
            r"\n\s*const legendItemsData = \[[\s\S]*?tempLegendSvg\.remove\(\);\s*\n",
            "\n",
            text2,
            count=1,
        )
        text2 = re.sub(
            r"\n\s*const legendRows = \[[\s\S]*?legendRows\.push\(currentRow\);\s*\n\s*\}\s*\n",
            "\n",
            text2,
            count=1,
        )
    return text2, n


def replace_multiple_pie_legend(text: str) -> tuple[str, int]:
    if "legendRows.forEach" not in text or "allGroups.forEach" not in text:
        return text, 0
    if "chartUtils.legend." in text:
        return text, 0
    pattern = (
        r"\n\s*// ---------- 5\.1 创建图例 ----------\s*\n"
        r"[\s\S]*?"
        r"\s*currentX \+= itemWidth;\s*\n"
        r"\s*\}\);\s*\n"
        r"\s*\}\);\s*\n"
    )
    replacement = (
        "\n    const allGroups = Array.from(new Set(chartData.map(d => d[groupField])));\n"
        "    const legendLayoutInfo = chartUtils.legend.layout(allGroups, {\n"
        "        colorResolver,\n"
        "        colors,\n"
        "        maxWidth: width - 40,\n"
        "        align: \"center\",\n"
        "        context: svg,\n"
        "        fontSize: typography.annotation.font_size,\n"
        "        fontFamily: typography.annotation.font_family,\n"
        "        fontWeight: typography.annotation.font_weight,\n"
        "    });\n"
        "    margin.top = legendLayoutInfo.height + 20;\n"
        "    chartUtils.legend.draw(svg, allGroups, {\n"
        "        colorResolver,\n"
        "        colors,\n"
        "    }, {\n"
        "        x: margin.left,\n"
        "        y: 10,\n"
        "        maxWidth: width - 40,\n"
        "        align: \"center\",\n"
        "        markerShape: \"circle\",\n"
        "        markerSize: 15,\n"
        "        labelGap: 5,\n"
        "        itemGap: 15,\n"
        "        rowGap: 12,\n"
        "        fontSize: typography.annotation.font_size,\n"
        "        fontFamily: typography.annotation.font_family,\n"
        "        fontWeight: typography.annotation.font_weight,\n"
        "        textColor: colorResolver.text({ fallback: \"#333333\" }).value,\n"
        "    });\n"
    )
    return subn(pattern, replacement, text, count=1)


def replace_stacked_bar_legend_block(text: str) -> tuple[str, int]:
    if "calculateLegendItemWidths" not in text:
        return text, 0
    pattern = (
        r"\n\s*// ---------- 1[34]\. 创建图例[\s\S]*?"
        r"(?:\s*tempSvg\.remove\(\);\s*\n)?"
        r"(?=\s*// ----------)"
    )
    replacement = (
        "\n    chartUtils.legend.draw(svg, groups, {\n"
        "        color: group => getColor(group),\n"
        "        colorResolver,\n"
        "        colors,\n"
        "    }, {\n"
        "        x: margin.left,\n"
        "        y: margin.top - 30,\n"
        "        maxWidth: barChartWidth,\n"
        "        fitToWidth: true,\n"
        "        minFontSize: 8,\n"
        "        markerShape: \"rect\",\n"
        "        markerSize: 12,\n"
        "        labelGap: 5,\n"
        "        itemGap: 15,\n"
        "        fontSize: typography.label.font_size,\n"
        "        fontFamily: typography.label.font_family,\n"
        "        fontWeight: typography.label.font_weight,\n"
        "        textColor: colorResolver.text({ fallback: \"#333333\" }).value,\n"
        "    });\n"
    )
    return subn(pattern, replacement, text, count=1)


def replace_proportional_rows_legend(text: str) -> tuple[str, int]:
    if "const rows = []" not in text or "uniqueGroupsForLegend.forEach" not in text:
        return text, 0
    if "legend-grad-" in text:
        return text, 0
    pattern = (
        r"\n\s*/\* ============ 8\. 创建图例 ============ \*/\s*\n"
        r"[\s\S]*?"
        r"(?=\s*return svg\.node\(\);)"
    )
    replacement = (
        "\n    const uniqueGroupsForLegend = uniqueGroups;\n"
        "    chartUtils.legend.draw(svg, uniqueGroupsForLegend, {\n"
        "        colorResolver,\n"
        "        colors,\n"
        "        colorScale,\n"
        "    }, {\n"
        "        x: margin.left,\n"
        "        y: 20,\n"
        "        maxWidth: W - 40,\n"
        "        align: \"center\",\n"
        "        markerShape: \"circle\",\n"
        "        markerSize: 12,\n"
        "        labelGap: 8,\n"
        "        itemGap: 20,\n"
        "        rowGap: 10,\n"
        "        fontSize: 12,\n"
        "        fontFamily: \"Arial\",\n"
        "    });\n"
    )
    return subn(pattern, replacement, text, count=1)


def replace_proportional_gradient_legend(text: str) -> tuple[str, int]:
    if "legend-grad-" not in text or "uniqueGroupsForLegend.forEach" not in text:
        return text, 0
    pattern = (
        r"\n\s*/\* ============ 8\. 创建图例 ============ \*/\s*\n"
        r"[\s\S]*?"
        r"\s*\.text\(group\);\s*\n"
        r"\s*\}\);\s*\n"
        r"(?=\s*return svg\.node\(\);)"
    )
    replacement = (
        "\n    const uniqueGroupsForLegend = [...new Set(raw.map(d => d[groupField]))];\n"
        "    const legendTitle = groupCol?.description || \"Categories\";\n"
        "    const legendTitleG = svg.append(\"g\")\n"
        "        .attr(\"class\", \"legend-title\")\n"
        "        .attr(\"transform\", `translate(${margin.left}, 20)`);\n"
        "    legendTitleG.append(\"text\")\n"
        "        .attr(\"x\", 0).attr(\"y\", 10)\n"
        "        .attr(\"font-size\", \"14px\").attr(\"font-weight\", \"bold\").attr(\"font-family\", \"Arial\")\n"
        "        .text(legendTitle);\n"
        "    const gradientLegendItems = chartUtils.legend.withGradientFills(\n"
        "        defs, uniqueGroupsForLegend, colorScale, colorResolver\n"
        "    );\n"
        "    chartUtils.legend.draw(svg, gradientLegendItems, {}, {\n"
        "        x: margin.left,\n"
        "        y: 30,\n"
        "        maxWidth: W,\n"
        "        align: \"center\",\n"
        "        markerShape: \"circle\",\n"
        "        markerSize: 16,\n"
        "        labelGap: 9,\n"
        "        itemGap: 15,\n"
        "        rowGap: 10,\n"
        "        fontSize: 12,\n"
        "        fontFamily: \"Arial\",\n"
        "    });\n"
    )
    return subn(pattern, replacement, text, count=1)


def replace_legend_group_offset(text: str) -> tuple[str, int]:
    if "chartUtils.legend." in text:
        return text, 0
    pattern = (
        r"\n\s*const legendGroup = svg\.append\(\"g\"\)\s*\n"
        r"\s*\.attr\(\"transform\", `translate\(\$\{margin\.left\}, \$\{margin\.top - 30\}\)`\);\s*\n"
        r"\s*let currentX = 0;\s*\n"
        r"\s*groups\.forEach\(\(group, i\) => \{\s*\n"
        r"\s*const legendItem = legendGroup\.append\(\"g\"\)\s*\n"
        r"[\s\S]*?"
        r"\s*currentX \+= legendItemWidths\[i\] \+ spacingBetweenItems;\s*\n"
        r"\s*\}\);\s*\n"
        r"\s*tempSvg\.remove\(\);\s*\n"
    )
    if "spacingBetweenItems" not in text:
        return text, 0
    replacement = (
        "\n    chartUtils.legend.draw(svg, groups, {\n"
        "        color: group => getColor(group),\n"
        "        colorResolver,\n"
        "        colors,\n"
        "    }, {\n"
        "        x: margin.left,\n"
        "        y: margin.top - 30,\n"
        "        maxWidth: barChartWidth,\n"
        "        fitToWidth: true,\n"
        "        minFontSize: 8,\n"
        "        markerShape: \"rect\",\n"
        "        markerSize: 12,\n"
        "        labelGap: 5,\n"
        "        itemGap: 15,\n"
        "        fontSize: typography.label.font_size,\n"
        "        fontFamily: typography.label.font_family,\n"
        "        fontWeight: typography.label.font_weight,\n"
        "        textColor: colorResolver.text({ fallback: \"#333333\" }).value,\n"
        "    });\n"
    )
    return subn(pattern, replacement, text, count=1)


def replace_vertical_bar_simple(text: str) -> tuple[str, int]:
    if "chartUtils.legend." in text:
        return text, 0
    pattern = (
        r"\n\s*//[^\n]*图例[^\n]*\n"
        r"\s*const legend = svg\.append\(\"g\"\)\s*\n"
        r"\s*\.attr\(\"class\", \"legend\"\)\s*\n"
        r"\s*\.attr\(\"transform\", `translate\(\$\{[^`]+\}, \$\{[^`]+\}\)`\);\s*\n"
        r"\s*(?:let )?legendOffset = 0;\s*\n"
        r"\s*\w+\.forEach\(\([^)]+\) => \{\s*\n"
        r"\s*const legendItem = legend\.append\(\"g\"\)\s*\n"
        r"[\s\S]*?"
        r"\s*legendOffset \+= [^;]+;\s*\n"
        r"\s*\}\);\s*\n"
    )
    match = re.search(pattern, text)
    if not match:
        return text, 0
    block = match.group(0)
    gm = re.search(r"(\w+)\.forEach\(\((\w+)", block)
    if not gm:
        return text, 0
    groups = gm.group(1)
    tm = re.search(r"translate\(\$\{([^}]+)\}, \$\{([^}]+)\}\)", block)
    if not tm:
        return text, 0
    cx, cy = tm.group(1).strip(), tm.group(2).strip()
    color = "colorScale,\n        colorResolver,\n        colors,"
    if "getColor(" in block:
        color = "color: d => getColor(d),\n        colorResolver,\n        colors,"
    replacement = (
        f"\n    chartUtils.legend.centered(svg, {groups}, {{\n"
        f"        {color}\n"
        f"    }}, {cx}, {cy}, {{\n"
        f"        markerShape: \"rect\",\n"
        f"        markerSize: 12,\n"
        f"        labelGap: 5,\n"
        f"        itemGap: 15,\n"
        f"        fontSize: typography.label.font_size,\n"
        f"        fontFamily: typography.label.font_family,\n"
        f"        fontWeight: typography.label.font_weight,\n"
        f"        textColor: colorResolver.text({{ fallback: \"#333333\" }}).value,\n"
        f"    }});\n"
    )
    return text[: match.start()] + replacement + text[match.end() :], 1


def replace_capsule_foreach(text: str) -> tuple[str, int]:
    if "legendCapsuleHeightH" not in text or "legendRows.forEach" in text:
        return text, 0
    if "chartUtils.legend.draw" in text and "capsule" in text:
        return text, 0
    pattern = (
        r"\n\s*// 布局图例项\s*\n"
        r"\s*let currentX = [\s\S]*?"
        r"\s*currentX \+= itemWidth \+ legendPadding;[^\n]*\n"
        r"\s*\}\);\s*\n"
    )
    if "legendCapsuleHeightH" not in text:
        return text, 0
    y_m = re.search(r"translate\(\$\{margin\.left\}, \$\{([^}]+)\}\)", text)
    y_expr = y_m.group(1).strip() if y_m else "margin.top * 0.7"
    replacement = (
        "\n    chartUtils.legend.draw(svg, groups, {\n"
        "        color: (group, index) => `url(#pattern-${group.replace(/[^a-zA-Z0-9]/g, '-')}-${index})`,\n"
        "        colorResolver,\n"
        "        colors,\n"
        "    }, {\n"
        f"        x: margin.left,\n"
        f"        y: {y_expr},\n"
        "        maxWidth: width - 100,\n"
        "        markerShape: \"capsule\",\n"
        "        markerWidth: legendCapsuleWidthH,\n"
        "        markerHeight: legendCapsuleHeightH,\n"
        "        labelGap: legendTextPadding,\n"
        "        itemGap: legendPadding,\n"
        "        rowGap: legendVerticalSpacing || 10,\n"
        "        fontSize: typography.label.font_size,\n"
        "        fontFamily: typography.label.font_family,\n"
        "        fontWeight: typography.label.font_weight,\n"
        "        textColor: colorResolver.text({ fallback: \"#333333\" }).value,\n"
        "    });\n"
    )
    text2, n = subn(pattern, replacement, text, count=1)
    if n:
        text2 = re.sub(
            r"\n\s*const tempLegendSvg = d3\.select\(containerSelector\)[\s\S]*?tempLegendSvg\.remove\(\);\s*\n",
            "\n",
            text2,
            count=1,
        )
        text2 = re.sub(
            r"\n\s*// 图例项样式参数[\s\S]*?const rowHeight = [^;]+;\s*\n",
            "\n",
            text2,
            count=1,
        )
    return text2, n


def replace_legend_other_centered(text: str) -> tuple[str, int]:
    pattern = (
        r"\n\s*const legend = svg\.append\(\"g\"\)\s*\n"
        r"\s*\.attr\(\"class\", \"other\"\)\s*\n"
        r"\s*\.attr\(\"transform\", `translate\(\$\{width / 2\}, \$\{margin\.top / 2\}\)`\);\s*\n"
        r"\s*let legendOffset = 0;\s*\n"
        r"\s*groups\.forEach\(\(group, i\) => \{\s*\n"
        r"[\s\S]*?"
        r"\s*legendOffset \+= legendItemWidths\[i\];\s*\n"
        r"\s*\}\);\s*\n"
        r"\s*legend\.attr\(\"transform\", `translate\(\$\{\(width - totalLegendWidth\) / 2\}, \$\{margin\.top / 2\}\)`\);\s*\n"
    )
    replacement = (
        "\n    chartUtils.legend.centered(svg, groups, {\n"
        "        colorScale,\n"
        "        colorResolver,\n"
        "        colors,\n"
        "    }, width / 2, margin.top / 2, {\n"
        "        markerShape: \"rect\",\n"
        "        markerSize: 15,\n"
        "        labelGap: 5,\n"
        "        itemGap: 10,\n"
        "        fontSize: typography.label.font_size,\n"
        "        fontFamily: typography.label.font_family,\n"
        "        fontWeight: typography.label.font_weight,\n"
        "        textColor: colorResolver.text({ fallback: \"#333333\" }).value,\n"
        "        markerRadius: variables.has_rounded_corners ? 2 : 0,\n"
        "    });\n"
    )
    text2, n = subn(pattern, replacement, text, count=1)
    if n:
        text2 = re.sub(
            r"\n\s*// 计算图例项宽度[\s\S]*?tempSvg\.remove\(\);\s*\n",
            "\n",
            text2,
            count=1,
        )
    return text2, n


def replace_display_groups_legend(text: str) -> tuple[str, int]:
    if "displayGroups.forEach" not in text or "legendItems.push" not in text:
        return text, 0
    pattern = (
        r"\n\s*//\s*7\.添加组图例\s*\n"
        r"[\s\S]*?"
        r"\s*// 如果仍然不适合[\s\S]*?\}\s*\n"
        r"\s*\}\s*\n"
    )
    replacement = (
        "\n    chartUtils.legend.draw(svg, displayGroups, {\n"
        "        color: group => getColor(group),\n"
        "        colorResolver,\n"
        "        colors,\n"
        "    }, {\n"
        "        x: margin.left,\n"
        "        y: margin.top - 20,\n"
        "        maxWidth: innerWidth + margin.right - 20,\n"
        "        fitToWidth: true,\n"
        "        minFontSize: 8,\n"
        "        markerShape: \"rect\",\n"
        "        markerSize: legendSquareSize,\n"
        "        labelGap: legendSpacing,\n"
        "        itemGap: legendGap,\n"
        "        fontSize: typography.label.font_size,\n"
        "        fontFamily: typography.label.font_family,\n"
        "        fontWeight: typography.label.font_weight,\n"
        "        textColor: colorResolver.text({ fallback: \"#333333\" }).value,\n"
        "    });\n"
    )
    return subn(pattern, replacement, text, count=1)


def replace_side_vertical_legend(text: str) -> tuple[str, int]:
    pattern = (
        r"\n\s*// Add legend\s*\n"
        r"\s*const legend = svg\.append\(\"g\"\)\s*\n"
        r"\s*\.attr\(\"transform\", `translate\(\$\{width - margin\.right \+ 20\}, \$\{margin\.top\}\)`\);\s*\n"
        r"\s*groupValues\.forEach\(\(group, i\) => \{\s*\n"
        r"\s*const legendRow = legend\.append\(\"g\"\)\s*\n"
        r"[\s\S]*?"
        r"\s*\}\);\s*\n"
    )
    replacement = (
        "\n    chartUtils.legend.draw(svg, groupValues, {\n"
        "        color: group => getColor(group),\n"
        "        colorResolver,\n"
        "        colors,\n"
        "    }, {\n"
        "        x: width - margin.right + 20,\n"
        "        y: margin.top,\n"
        "        direction: \"vertical\",\n"
        "        markerShape: \"rect\",\n"
        "        markerSize: 15,\n"
        "        labelGap: 5,\n"
        "        rowGap: 10,\n"
        "        itemHeight: 25,\n"
        "        fontSize: typography.label.font_size,\n"
        "        fontFamily: typography.label.font_family,\n"
        "        fontWeight: typography.label.font_weight,\n"
        "        textColor: colorResolver.text({ fallback: \"#333333\" }).value,\n"
        "        markerRadius: variables.has_rounded_corners ? 2 : 0,\n"
        "    });\n"
    )
    return subn(pattern, replacement, text, count=1)


def migrate_file(path: Path) -> bool:
    original = path.read_text(encoding="utf-8")
    if "chartUtils.legend." in original and "legendItem.append" not in original:
        return False
    text = original
    for fn in (
        replace_capsule_rows,
        replace_capsule_foreach,
        replace_multiple_pie_legend,
        replace_stacked_bar_legend_block,
        replace_legend_group_offset,
        replace_proportional_rows_legend,
        replace_proportional_gradient_legend,
        replace_vertical_bar_simple,
        replace_legend_other_centered,
        replace_display_groups_legend,
        replace_side_vertical_legend,
    ):
        text, _ = fn(text)
    if text != original:
        path.write_text(text, encoding="utf-8")
        return True
    return False


def main() -> None:
    changed = 0
    for path in sorted(TEMPLATE_ROOT.rglob("*.js")):
        if migrate_file(path):
            changed += 1
    print(f"pass2 migrated {changed} files")


if __name__ == "__main__":
    main()