#!/usr/bin/env python3
"""Normalize markdown heading styles to consistent ATX format."""
import re
import sys

_ATX_RE = re.compile(r'^(#{1,})(\s+)(.+)$')

def normalize(text: str) -> str:
    lines = text.splitlines()
    out = []
    i = 0
    while i < len(lines):
        line = lines[i]
        stripped = line.rstrip()
        next_line = lines[i + 1].rstrip() if i + 1 < len(lines) else ''
        # Setext h1/h2: only if current line is NOT already an ATX heading
        if not _ATX_RE.match(stripped):
            if re.fullmatch(r'=+\s*$', next_line):
                out.append(f'# {stripped}')
                i += 2
                continue
            if re.fullmatch(r'-+\s*$', next_line):
                out.append(f'## {stripped}')
                i += 2
                continue
        # ATX heading: cap level at 6
        m = _ATX_RE.match(stripped)
        if m:
            hashes, space, content = m.group(1), m.group(2), m.group(3)
            level = min(len(hashes), 6)
            out.append(f'{"#" * level}{space}{content}')
            i += 1
            continue
        out.append(stripped)
        i += 1
    return '\n'.join(out) + ('\n' if text.endswith('\n') else '')

def main() -> int:
    if len(sys.argv) < 2:
        print('usage: md-heading-normalize.py <file>', file=sys.stderr)
        return 2
    path = sys.argv[1]
    with open(path, 'r', encoding='utf-8') as f:
        src = f.read()
    sys.stdout.write(normalize(src))
    return 0

if __name__ == '__main__':
    raise SystemExit(main())