"""Symbolic compile-repair for COBOL SML drafts.

The tiny model gets COBOL structure ~right but drops mechanical tokens (a `)` in
`PIC 9(4)`, the name after `PROGRAM-ID.`, emits dup lines). Compilation is
all-or-nothing, so these kill it. This is the "lean into symbolic" layer: fix
MECHANICS only (parens, program name, dedup, END PROGRAM) — never invent
PROCEDURE logic. Honest: repair is bounded structural surgery, not generation.

repair(src) -> repaired_src. Idempotent-ish; safe to run then re-compile.
"""
import re

NAME_RE = re.compile(r"[A-Z][A-Z0-9-]*", re.I)


def _program_name(lines):
    for pat in (r"END\s+PROGRAM\s+([A-Za-z][A-Za-z0-9-]*)",
                r"PROGRAM-ID\.\s*([A-Za-z][A-Za-z0-9-]*)"):
        for l in lines:
            m = re.search(pat, l, re.I)
            if m:
                return m.group(1)
    return "PROG1"


def _balance_parens(line):
    """Close unbalanced '(' (the `PIC 9(4.` failure). Insert before trailing period."""
    opens, closes = line.count("("), line.count(")")
    if opens > closes:
        miss = ")" * (opens - closes)
        s = line.rstrip()
        return (s[:-1] + miss + ".") if s.endswith(".") else (s + miss)
    return line


DIV_SEC_RE = re.compile(r"^\s*[A-Z-]+\s+(DIVISION|SECTION)\b.*$", re.I)


def repair(src: str) -> str:
    lines = [l for l in src.splitlines() if l.strip() != ""]

    # 1) drop consecutive duplicate lines (dup PROGRAM-ID / OCCURS the model emits)
    dedup = []
    for l in lines:
        if not dedup or dedup[-1].strip() != l.strip():
            dedup.append(l)
    lines = dedup

    # 1b) OCCURS n PIC -> OCCURS n TIMES PIC (model drops TIMES)
    lines = [re.sub(r"\bOCCURS\s+(\d+)\s+(PIC)\b", r"OCCURS \1 TIMES \2", l, flags=re.I)
             for l in lines]

    # 1c) DIVISION/SECTION headers must end with a period
    lines = [l.rstrip() + ("." if DIV_SEC_RE.match(l) and not l.rstrip().endswith(".") else "")
             for l in lines]

    # 1d) drop a duplicate data field declaration by (level, name) when a later
    #     well-formed one exists — the model sometimes emits a malformed twin.
    seen = {}
    field = re.compile(r"^\s*(\d\d)\s+([A-Za-z][A-Za-z0-9-]*)\b")
    for i, l in enumerate(lines):
        m = field.match(l)
        if m:
            seen.setdefault(m.groups(), []).append(i)
    drop = set()
    for _, idxs in seen.items():
        if len(idxs) > 1:
            # keep the one that has PIC or OCCURS...TIMES; drop the others
            best = max(idxs, key=lambda j: ("PIC" in lines[j].upper()) + ("TIMES" in lines[j].upper()))
            for j in idxs:
                if j != best:
                    drop.add(j)
    lines = [l for i, l in enumerate(lines) if i not in drop]

    name = _program_name(lines)

    # 2) exactly one `PROGRAM-ID. <name>.` right after IDENTIFICATION DIVISION
    lines = [l for l in lines if not re.match(r"\s*PROGRAM-ID", l, re.I)]
    if not any(re.match(r"\s*IDENTIFICATION", l, re.I) for l in lines):
        lines = ["IDENTIFICATION DIVISION."] + lines
    out = []
    for l in lines:
        out.append(l)
        if re.match(r"\s*IDENTIFICATION\s+DIVISION", l, re.I):
            out.append(f"PROGRAM-ID. {name}.")
    lines = out

    # 3) per-line paren balance (fixes `PIC 9(4.`)
    lines = [_balance_parens(l) for l in lines]

    # 4) ensure a single well-formed END PROGRAM <name>.
    lines = [l for l in lines if not re.search(r"END\s+PROGRAM", l, re.I)]
    lines.append(f"END PROGRAM {name}.")

    return "\n".join(lines)


if __name__ == "__main__":
    broken = """IDENTIFICATION DIVISION.
DATA DIVISION.
WORKING-STORAGE SECTION.
01 WS-I PIC 9(4.
LINKAGE SECTION.
01 LINKED-ITEMS.
05 L-TAB OCCURS 12 TIMES PIC S9(5).
05 RESULT PIC S9(9).
PROCEDURE DIVISION USING LINKED-ITEMS.
MOVE L-TAB(1) TO RESULT
PERFORM VARYING WS-I FROM 2 BY 1 UNTIL WS-I > 5
IF L-TAB(WS-I) < RESULT
MOVE L-TAB(WS-I) TO RESULT
END-IF
END-PERFORM
GOBACK.
END PROGRAM."""
    print(repair(broken))