#!/usr/bin/env python3
"""Validate Mermaid diagrams in markdown files."""

import glob
import re
import sys


# Diagram types that are valid in Mermaid
VALID_TYPES = {
    "flowchart",
    "graph",
    "classDiagram",
    "sequenceDiagram",
    "stateDiagram",
    "stateDiagram-v2",
    "erDiagram",
    "gantt",
    "pie",
    "mindmap",
    "journey",
    "gitGraph",
    "requirementDiagram",
}


def extract_mermaid_blocks(content: str) -> list[tuple[int, str]]:
    """Extract mermaid blocks with their positions."""
    blocks = []
    for match in re.finditer(r"```mermaid\n(.*?)\n```", content, re.DOTALL):
        blocks.append((match.start(), match.group(1)))
    return blocks


def validate_block(block: str) -> bool:
    """Check if a mermaid block has a valid diagram type."""
    if not block.strip():
        return True  # Empty block is OK
    first_line = block.strip().split("\n")[0]
    return any(first_line.startswith(t) for t in VALID_TYPES)


def main():
    md_files = glob.glob("docs/*.md")

    errors = []
    for filepath in md_files:
        content = open(filepath).read()
        blocks = extract_mermaid_blocks(content)

        for i, (_, block) in enumerate(blocks):
            if not validate_block(block):
                errors.append(f"{filepath}: invalid diagram type in block {i + 1}")

    if errors:
        for e in errors:
            print(f"ERROR: {e}")
        sys.exit(1)

    print(f"Validated {len(md_files)} markdown files - all OK")


if __name__ == "__main__":
    main()