#!/usr/bin/env python3
"""Add :lc-problem <number> to org source blocks in leetcode problem files."""

import re
import sys
from pathlib import Path


def process_file(path: Path, dry_run: bool = False) -> bool:
    """Add :lc-problem to src blocks that don't already have it. Returns True if modified."""
    m = re.match(r"(\d{4})-", path.stem)
    if not m:
        return False
    problem_num = str(int(m.group(1)))  # strip leading zeros: "0015" -> "15"

    text = path.read_text()
    original = text

    # Match #+begin_src <lang> where lang is python or cpp
    def fix_src_block(match: re.Match) -> str:
        full = match.group(0)
        lang = match.group(1)
        suffix = match.group(2) or ""

        # If already has :lc-problem, fix leading zeros
        if ":lc-problem" in suffix:
            suffix = re.sub(r":lc-problem 0+(\d)", rf":lc-problem \1", suffix)
            return f"#+begin_src {lang}{suffix}"

        # Add :lc-problem
        extra = f" :lc-problem {problem_num}"
        if lang == "python":
            extra += " :lc-lang python3"
        return f"#+begin_src {lang}{extra}{suffix}"

    text = re.sub(
        r"#\+begin_src (python|cpp)([^\n]*)",
        fix_src_block,
        text,
    )

    if text != original:
        if dry_run:
            print(f"  would modify: {path.name}")
        else:
            path.write_text(text)
            print(f"  modified: {path.name}")
        return True
    return False


def main():
    dry_run = "--dry-run" in sys.argv
    base = Path(__file__).parent / "dsa"
    if not base.exists():
        print(f"dsa/ directory not found at {base}")
        sys.exit(1)

    org_files = sorted(base.rglob("*.org"))
    print(f"Found {len(org_files)} org files in dsa/")

    modified = 0
    skipped = 0
    for f in org_files:
        if process_file(f, dry_run):
            modified += 1
        else:
            skipped += 1

    print(f"\nDone: {modified} modified, {skipped} unchanged")


if __name__ == "__main__":
    main()