import re
from django import template
from django.utils.text import slugify
from django.utils.safestring import mark_safe

register = template.Library()


@register.simple_tag
def get_toc(body):
    """Extract headings from StreamField body and return table of contents."""
    toc = []

    for block in body:
        if block.block_type == 'paragraph':
            # Extract h2 and h3 from rich text
            html = str(block.value)
            headings = re.findall(r'<(h[23])>(.*?)</\1>', html, re.IGNORECASE)

            for tag, text in headings:
                # Clean text (remove HTML tags)
                clean_text = re.sub(r'<[^>]+>', '', text)
                slug = slugify(clean_text)
                level = int(tag[1])  # 2 or 3
                toc.append({
                    'text': clean_text,
                    'slug': slug,
                    'level': level,
                })

    return toc


@register.filter
def add_heading_ids(html):
    """Add id attributes to h2 and h3 tags for anchor links."""
    def replace_heading(match):
        tag = match.group(1)
        content = match.group(2)
        clean_text = re.sub(r'<[^>]+>', '', content)
        slug = slugify(clean_text)
        return f'<{tag} id="{slug}">{content}</{tag}>'

    result = re.sub(r'<(h[23])>(.*?)</\1>', replace_heading, str(html), flags=re.IGNORECASE)
    return mark_safe(result)
