"""COBOL-aware word/symbol tokenizer for the COBOL SML.

Same spirit as TinkyBrain's word-level tokenizer (the vocab IS the boundary),
but COBOL-shaped:
  * case PRESERVED (COBOL identifiers matter; keywords are case-insensitive but
    the corpus is uppercase, so preserving keeps the vocab tight)
  * hyphenated names stay whole (PROGRAM-ID, L-TAB, END-PERFORM)
  * string literals "..." / '...' are single tokens
  * numbers are single tokens
  * every structural symbol ( ) . , : ; etc. is its own token
  * <nl> encodes line structure so free-format COBOL round-trips

Free-format is the target dialect: strip fixed-format indentation, compile with
`cobc -free -c`. Proven equivalent to the fixed-format source (see de-risk test).

The detokenizer re-applies COBOL spacing (no space around parens / before . , )
so `PIC S9(5)`, `L-TAB(WS-I)`, `RESULT.` reconstruct exactly.
"""

import json
import re
from collections import Counter

# Special tokens (mirror TinkyBrain ids so muscle memory carries over)
PAD, BOS, EOS, SEP, UNK, NL = 0, 1, 2, 3, 4, 5
SPECIAL = {"<pad>": PAD, "<bos>": BOS, "<eos>": EOS, "<sep>": SEP, "<unk>": UNK, "<nl>": NL}

# One token per: string literal | hyphenated word/keyword | number | single symbol
_TOKEN_RE = re.compile(r"""
      "[^"]*"            # double-quoted literal
    | '[^']*'            # single-quoted literal
    | [A-Za-z][A-Za-z0-9-]*   # identifier / keyword (hyphens kept)
    | [0-9]+             # integer
    | [^\sA-Za-z0-9]     # any single structural symbol
""", re.VERBOSE)

# Detokenizer spacing rules
_NO_SPACE_BEFORE = {".", ",", ")", "(", ";", ":"}
_NO_SPACE_AFTER = {"("}


def tokenize_cobol(src: str):
    """COBOL source -> flat token list with <nl> between lines (indentation dropped)."""
    out = []
    for i, line in enumerate(src.splitlines()):
        if i > 0:
            out.append("<nl>")
        out.extend(_TOKEN_RE.findall(line.strip()))
    return out


def detokenize_cobol(tokens) -> str:
    """Token list -> free-format COBOL source (re-applies COBOL spacing)."""
    parts, prev = [], None
    for t in tokens:
        if t == "<nl>":
            parts.append("\n")
            prev = None
            continue
        if prev is None or prev == "\n":
            parts.append(t)
        elif prev in _NO_SPACE_AFTER or t in _NO_SPACE_BEFORE:
            parts.append(t)
        else:
            parts.append(" " + t)
        prev = t
    return "".join(parts)


class CobolTokenizer:
    def __init__(self, vocab_size=4096):
        self.vocab_size = vocab_size
        self.word2idx = dict(SPECIAL)
        self.idx2word = {v: k for k, v in SPECIAL.items()}
        self.fitted = False

    def fit(self, texts):
        counts = Counter()
        for t in texts:
            counts.update(tok for tok in _TOKEN_RE.findall(t) if tok)  # NL handled separately
        for word, _ in counts.most_common(self.vocab_size - len(SPECIAL)):
            if len(self.word2idx) >= self.vocab_size:
                break
            if word not in self.word2idx:
                self.word2idx[word] = len(self.word2idx)
                self.idx2word[len(self.idx2word)] = word
        self.fitted = True
        print(f"CobolTokenizer: {len(self.word2idx)} tokens (from {len(counts)} unique)")

    def encode(self, text: str):
        return [self.word2idx.get("<nl>", NL) if t == "<nl>" else self.word2idx.get(t, UNK)
                for t in tokenize_cobol(text)]

    def encode_tokens(self, tokens):
        return [self.word2idx.get(t, UNK) for t in tokens]

    def decode(self, ids):
        toks = [self.idx2word.get(i, "<unk>") for i in ids if i not in (PAD, BOS, EOS, SEP)]
        return detokenize_cobol(toks)

    def save(self, path):
        json.dump(self.word2idx, open(path, "w"))

    @classmethod
    def load(cls, path):
        w2i = json.load(open(path))
        t = cls(vocab_size=len(w2i))
        t.word2idx = {k: int(v) for k, v in w2i.items()}
        t.idx2word = {int(v): k for k, v in w2i.items()}
        t.fitted = True
        return t