#!/usr/bin/env python3
"""Render the wiki's accessible SVG diagrams using only Python's standard library."""

import html
import json
import re
import textwrap
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
RAW = "https://raw.githubusercontent.com/wiki/ModelCloud/GPTQModel/assets/diagrams"
COLORS = [
    ("#1d4ed8", "#eff6ff", "#bfdbfe"),
    ("#0f766e", "#f0fdfa", "#99f6e4"),
    ("#7e22ce", "#faf5ff", "#e9d5ff"),
    ("#9a3412", "#fff7ed", "#fed7aa"),
]


def wrapped(text, width):
    return textwrap.wrap(text, width=width, break_long_words=True, break_on_hyphens=False)


def text(x, y, lines, size=18, color="#334155", weight="400", leading=25):
    content = "".join(
        f'<tspan x="{x}" dy="{0 if i == 0 else leading}">{html.escape(line)}</tspan>' for i, line in enumerate(lines)
    )
    return f'<text x="{x}" y="{y}" font-size="{size}" font-weight="{weight}" fill="{color}">{content}</text>'


def render(diagram):
    title = wrapped(diagram["title"], 53)
    subtitle = wrapped(diagram["subtitle"], 90)
    header = 58 + len(title) * 36 + len(subtitle) * 26 + 28
    note = wrapped(diagram["note"], 92)
    lane_heights = []
    for lane in diagram["lanes"]:
        count = len(lane["steps"])
        assert 2 <= count <= 4
        width = (920 - (count - 1) * 32) / count
        max_title = max(len(wrapped(step["title"], int((width - 34) / 13))) for step in lane["steps"])
        max_detail = max(len(wrapped(step["detail"], int((width - 34) / 11))) for step in lane["steps"])
        lane_heights.append(34 + max_title * 26 + max_detail * 25 + 60)
    height = header + sum(lane_heights) + len(note) * 25 + 68
    description = (
        diagram["subtitle"]
        + " "
        + " ".join(
            lane["label"] + ": " + " Then ".join(step["title"] + ". " + step["detail"] for step in lane["steps"])
            for lane in diagram["lanes"]
        )
        + " "
        + diagram["note"]
    )
    parts = [
        '<?xml version="1.0" encoding="UTF-8"?>',
        f'<svg xmlns="http://www.w3.org/2000/svg" width="1000" height="{height}" '
        f'viewBox="0 0 1000 {height}" role="img" aria-labelledby="title description">',
        f'<title id="title">{html.escape(diagram["title"])}</title>',
        f'<desc id="description">{html.escape(description)}</desc>',
        '<defs><marker id="arrow" markerWidth="8" markerHeight="8" refX="6" refY="4" '
        'orient="auto"><path d="M0 1 L6 4 L0 7" fill="none" stroke="#64748b" '
        'stroke-width="1.5"/></marker></defs>',
        f'<rect x="1" y="1" width="998" height="{height - 2}" rx="16" fill="#ffffff" stroke="#cbd5e1"/>',
        '<g font-family="DejaVu Sans, Arial, sans-serif">',
        text(40, 34, ["GPTQMODEL  /  WORKFLOW"], 13, "#475569", "700"),
        text(40, 76, title, 28, "#0f172a", "700", 36),
        text(40, 76 + len(title) * 36, subtitle, 18, "#475569", leading=26),
    ]
    y = header
    for row, (lane, lane_height) in enumerate(zip(diagram["lanes"], lane_heights)):
        color, fill, border = COLORS[row % len(COLORS)]
        parts.append(text(40, y, [lane["label"]], 19, color, "700"))
        count = len(lane["steps"])
        width = (920 - (count - 1) * 32) / count
        box_y = y + 16
        box_height = lane_height - 62
        for col, step in enumerate(lane["steps"]):
            x = 40 + col * (width + 32)
            step_title = wrapped(step["title"], int((width - 34) / 13))
            step_detail = wrapped(step["detail"], int((width - 34) / 11))
            parts.append(
                f'<rect x="{x}" y="{box_y}" width="{width}" height="{box_height}" '
                f'rx="10" fill="{fill}" stroke="{border}"/>'
            )
            parts.append(text(x + 16, box_y + 32, step_title, 20, "#0f172a", "700", 26))
            parts.append(text(x + 16, box_y + 45 + len(step_title) * 26, step_detail))
            if col < count - 1:
                cy = box_y + box_height / 2
                parts.append(
                    f'<path d="M{x + width + 5} {cy} H{x + width + 25}" '
                    'stroke="#64748b" stroke-width="2" marker-end="url(#arrow)"/>'
                )
        y += lane_height
    parts.append(f'<path d="M40 {y - 4} H960" stroke="#e2e8f0"/>')
    parts.append(text(40, y + 26, note, 17, "#475569", leading=25))
    parts.append("</g></svg>\n")
    return "\n".join(parts)


def markdown(diagram):
    slug = diagram["page"]
    asset = slug.lower() + ".svg"
    lines = [
        "<!-- workflow:start -->",
        "## Workflow at a glance",
        "",
        f"![{diagram['title']}. {diagram['subtitle']}]({RAW}/{asset})",
        "",
        f"**Reading the diagram:** {diagram['note']}",
        "",
        "<details>",
        "<summary>Text version of the workflow</summary>",
        "",
    ]
    for lane in diagram["lanes"]:
        steps = " → ".join(step["title"] for step in lane["steps"])
        lines.append(f"- **{lane['label']}:** {steps}.")
    lines += ["", "</details>", "", f"[Open full-size SVG]({RAW}/{asset})", "<!-- workflow:end -->"]
    return "\n".join(lines) + "\n"


def main():
    diagrams = json.loads((ROOT / ".tools/diagrams.json").read_text())
    assert len({d["page"] for d in diagrams}) == len(diagrams)
    for diagram in diagrams:
        assert re.fullmatch(r"[A-Za-z0-9-]+", diagram["page"])
        svg = ROOT / "assets/diagrams" / (diagram["page"].lower() + ".svg")
        svg.write_text(render(diagram))
        page = ROOT / (diagram["page"] + ".md")
        content = page.read_text()
        block = markdown(diagram)
        if "<!-- workflow:start -->" in content:
            content = re.sub(
                r"<!-- workflow:start -->.*?<!-- workflow:end -->\n?", lambda _: block, content, flags=re.S
            )
        else:
            # Keep the opening explanation before the visual; insert before the next section.
            matches = list(re.finditer(r"^## ", content, re.M))
            index = 1 if matches and content[matches[0].start() :].startswith("## Overview\n") else 0
            offset = matches[index].start() if len(matches) > index else len(content)
            content = content[:offset].rstrip() + "\n\n" + block + "\n" + content[offset:]
        page.write_text(content)
    print(f"Rendered and embedded {len(diagrams)} workflow diagrams.")


if __name__ == "__main__":
    main()