"""Shared word tokenization, actual SentencePiece special IDs, complete windows."""
import re

TOKEN_RE = re.compile(r'\w+(?:-\w+)*|[^\w\s]', re.UNICODE)
MAX_LEN = 256


def tokenizer_contract(sp):
    return dict(bos_token_id=sp.bos_id(), eos_token_id=sp.eos_id(),
                pad_token_id=sp.pad_id(), unk_token_id=sp.unk_id())


def word_windows(tokens, sp, max_len=MAX_LEN, overlap=0):
    """Yield (start_word, end_word, IDs, first-subword positions) without dropping words."""
    if max_len < 3 or overlap < 0:
        raise ValueError('Invalid window settings')
    pieces = [sp.encode(t, out_type=int) or [sp.unk_id()] for t in tokens]
    start = 0
    while start < len(tokens):
        ids, positions, end = [sp.bos_id()], [], start
        while end < len(tokens) and len(ids)+len(pieces[end])+1 <= max_len:
            positions.append(len(ids))
            ids.extend(pieces[end])
            end += 1
        if end == start:
            raise ValueError(f'Word at index {start} cannot fit into a {max_len}-subtoken window')
        ids.append(sp.eos_id())
        yield start, end, ids, positions
        if end == len(tokens):
            break
        start = max(start+1, end-overlap)


"""Body appended to the shared tokenization helpers when packaging the release."""
import hashlib
import json
from pathlib import Path
import numpy as np
import onnxruntime as ort
import sentencepiece as spm


class BashkirTokenLID:
    """Standalone CPU word-level LID with offsets and complete long-text windows."""
    def __init__(self, model_dir=None, precision='int8'):
        self.directory=Path(model_dir) if model_dir is not None else Path(__file__).resolve().parent
        config=json.loads((self.directory/'config.json').read_text(encoding='utf-8'))
        meta=json.loads((self.directory/'META.json').read_text(encoding='utf-8'))
        if precision not in config['models']:
            raise ValueError(f"precision must be one of {sorted(config['models'])}")
        model_name=config['models'][precision];tokenizer_name=config['tokenizer_file']
        for name in [model_name,tokenizer_name,'config.json']:
            with (self.directory/name).open('rb') as stream:
                digest=hashlib.sha256()
                for block in iter(lambda:stream.read(1024*1024),b''): digest.update(block)
                actual=digest.hexdigest()
            if actual!=meta['artifacts'][name]['sha256']:
                raise ValueError(f'Artifact checksum mismatch: {name}')
        self.tags=tuple(config['labels'])
        self.max_length=int(config['max_position_embeddings'])
        self.overlap=int(config['word_window_overlap'])
        self.sp=spm.SentencePieceProcessor(model_file=str(self.directory/tokenizer_name))
        expected={k:int(config[k]) for k in ['bos_token_id','eos_token_id','pad_token_id','unk_token_id']}
        if tokenizer_contract(self.sp)!=expected or self.sp.get_piece_size()!=config['vocab_size']:
            raise ValueError('Tokenizer does not match the model contract')
        options=ort.SessionOptions();options.intra_op_num_threads=4;options.inter_op_num_threads=1
        self.session=ort.InferenceSession(str(self.directory/model_name),sess_options=options,providers=['CPUExecutionProvider'])

    def predict(self,text):
        """Return token, label, model confidence, start and end for every token.

        Offsets index Python Unicode strings. Confidence is an uncalibrated model
        score. Empty text returns []; overlong individual words raise ValueError.
        """
        if not isinstance(text,str): raise TypeError('text must be a string')
        matches=list(TOKEN_RE.finditer(text));tokens=[m.group() for m in matches]
        result=[None]*len(tokens);priorities=[-1]*len(tokens)
        for start,end,ids,positions in word_windows(tokens,self.sp,max_len=self.max_length,overlap=self.overlap):
            logits=self.session.run(['logits'],{'input_ids':np.asarray([ids],dtype=np.int64)})[0][0,positions]
            probs=np.exp(logits-logits.max(-1,keepdims=True));probs/=probs.sum(-1,keepdims=True)
            for local,index in enumerate(range(start,end)):
                priority=min(local+1,end-start-local)
                if priority>priorities[index]:
                    tag=int(probs[local].argmax())
                    result[index]=dict(token=tokens[index],label=self.tags[tag],confidence=float(probs[local,tag]),
                        start=matches[index].start(),end=matches[index].end())
                    priorities[index]=priority
        if any(r is None for r in result): raise AssertionError('Missing token after window inference')
        return result


if __name__=='__main__':
    import argparse
    import sys
    if hasattr(sys.stdout,'reconfigure'): sys.stdout.reconfigure(encoding='utf-8')
    parser=argparse.ArgumentParser(description='Word-level Bashkir / Russian / Tatar LID')
    parser.add_argument('text');parser.add_argument('--precision',choices=['int8','fp32'],default='int8')
    args=parser.parse_args()
    print(json.dumps(BashkirTokenLID(precision=args.precision).predict(args.text),ensure_ascii=False,indent=2))