# python3 -I 06_ref.py <out.json>: my own transcription (from reading) of the official text functions, run on adversarial
# inputs. The official files are not imported. Output is compared with the platform's TS by 06_cmp.mts.
import re, json, random, sys, string, unicodedata
def remove_inner_thoughts(d):
    c=re.sub(r'\[.*?\]','',d)
    c='\n'.join(l.strip() for l in c.split('\n'))
    c=re.sub(r'\n+','\n',c)
    c=re.sub(r'\s*<\s*system[_\s]+think(?:ing)?\s*>.*?</\s*system[_\s]+think(?:ing)?\s*>\s*','',c,flags=re.S).strip()
    c=re.sub(r'\s*<\s*role[_\s]+think(?:ing)?\s*>.*?</\s*role[_\s]+think(?:ing)?\s*>\s*','',c,flags=re.S).strip()
    c=re.sub(r'\s*<\s*[_\s]+think(?:ing)?\s*>.*?</\s*[_\s]+think(?:ing)?\s*>\s*','',c,flags=re.S).strip()
    return c.strip()
def add_speaker_name(d,s):
    if any(l.strip().startswith(f"{s}:") or l.strip().startswith(f"{s}：") for l in d.split('\n')): return d
    return f"{s}: {d}"
def extract_nsp(r):
    m=re.findall(r'(?:\*\*|#)?\s*Next\s+Speaker\s*(?:\*\*|#)?\s*(?:[:\-]|\bis\b)?\s*\*?([^\n\*]+)',r,re.IGNORECASE)
    if m:
        c=m[-1].strip().rstrip('.,!"\'').replace('*','')
        if c: return c
    return None
def parse_nsp(r):
    n=extract_nsp(r)
    if n is None: n=r.split(':')[0].strip() if ':' in r else r
    return n
def _extract(text):
    text=re.sub(r'"([^"\\]*(\\.[^"\\]*)*)"',lambda m:m.group().replace('\n',r'\\n'),text)
    def safe(text):
        try: return json.loads(text)
        except json.JSONDecodeError:
            res=[];st=0
            while st<len(text):
                try:
                    o,e=json.JSONDecoder().raw_decode(text[st:]); res.append(o); st+=e
                except json.JSONDecodeError: st+=1
            return max(res,key=lambda x:len(json.dumps(x))) if res else None
    v=safe(text)
    return v if v else None
DIM=['Storyline Consistency','Anthropomorphism','Character Fidelity','Storyline Quality']
def parse_resp(r):
    try:
        assert isinstance(r,dict)
        for k,v in r.items():
            assert k in DIM
            assert 'flaws' in v
            for f in v['flaws']:
                if f.get('severity',None) is None: f['severity']=1
        return r
    except: return False
def read(text,dim,R):
    v=_extract(text)
    if v is None: return {'s':'fix'}
    p=parse_resp(v)
    if not p: return {'s':'retry'}
    if dim not in p or not isinstance(p[dim],dict): return {'s':'missing'}
    fl=p[dim]['flaws']
    try:
        S=sum([f['severity'] for f in fl if isinstance(f['severity'],int)])
    except Exception as e:
        return {'s':'crash'}
    return {'s':'ok','S':str(S),'score':max(0,min(100-(S-0.3*R)*5,100))}
def own(resp,name):
    cur=None;u="";out=[]
    for line in resp.split('\n'):
        if ':' in line:
            ch=line.split(':',1)[0].strip()
            if cur!=ch:
                if u: out.append((cur,u))
                cur=ch;u=""
        u+=line+"\n"
    if u: out.append((cur,u))
    x=[t for c,t in out if c==name]
    return x[0] if x else None
# tokenize_words with unicode classes emulated for ascii+cjk input (regex module is not installed here)
PS=''.join(ch for ch in map(chr,range(32,127)) if unicodedata.category(ch)[0] in 'PS')+'，。！？“”…'
TOK=re.compile(r'\b\w+\b|[一-鿿]|[぀-ゟ゠-ヿ]|\d|['+re.escape(PS)+']')
def tokenize(t):
    o=[]
    for tok in TOK.findall(t):
        if re.match(r'[一-鿿]|[぀-ゟ゠-ヿ]',tok): o.extend(list(tok))
        else: o.append(tok)
    return o
def fix(resp):
    def d1(tk,mn=5,mx=30):
        n=len(tk)
        for L in range(mn,min(mx+1,n+1)):
            for i in range(n-L+1):
                s=tk[i:i+L]; rep=True
                for r in range(1,5):
                    cp=i+r*L
                    if tk[cp:cp+L]!=s: rep=False;break
                if rep: return tk[:i+L]
        return False
    def d2(tk,mn=15,mx=30):
        n=len(tk);fr=999999999999999;fs={}
        for L in range(mn,min(mx+1,n+1)):
            c={}
            for i in range(n-L+1):
                s=tuple(tk[i:i+L])
                if c.get(s,0)>0:
                    if i-fs[s]>=L: fr=min(fr,i)
                else: fs[s]=i
                c[s]=c.get(s,0)+1
        return tk[:fr] if fr<999999999999999 else False
    def cat(tk):
        t="";last=None
        for k in tk:
            cu='CJK' if re.match(r'[一-鿿]|[぀-ゟ゠-ヿ]',k) else 'Other'
            if k in string.punctuation: cu='P'
            if last in['Other','P'] and cu=='Other': t+=" "+k
            else: t+=k
            last=cu
        if re.match(r'[a-zA-Z0-9]+$',t[-1]): t+='.'
        if re.match(r'[一-鿿]+$',t[-1]): t+='。'
        if re.match(r'[぀-ゟ゠-ヿ]+$',t[-1]): t+='。'
        return t
    sign=False
    ls=re.findall(r'[a-zA-Z]{100,}',resp)
    if ls:
        for s in ls: resp=resp.replace(s,s[:20])
        sign=True
    tk=tokenize(resp); x=d1(tk)
    if x==False: x=d2(tk)
    if x:
        resp=cat(x); sign=True
    return resp if sign else False
rnd=random.Random(20261101)
TH=['[a thought]','[multi\nline]','(acts)','Hello there.','  spaced  ','\n\n','\n','<system_think>x</system_think>','<role think>y\nz</role think>','< _thinking>q</ _thinking>','<think>k</think>','<_think>w</_think>','[]','[[nested]]',']','[','Alice:','Alice：','Bob:','  Alice: hi','：','　',' ','\x85','
','é','中文','Environment: rain falls.','\t']
NS=['Next Speaker: Alice','**Next Speaker:** Alice','**Next Speaker**: Alice','2. Next Speaker: Alice.','### Next Speaker\nAlice','Next speaker is Bob','next  speaker - Environment','Next Speaker: <END CHAT>','Next Speaker: "Alice"','1. Reasoning: The next speaker should be Bob because.\n2. Next Speaker: Alice','Alice','Alice: hello','<END CHAT>','random','Next Speaker:','Next Speaker: **Alice**','Next Speaker:\n\nAlice','Reasoning: x\nNext Speaker: Alice, the queen','NEXT SPEAKER: Sir Bob!','Next Speaker - *Alice*','Nextspeaker: Alice','#Next Speaker# Alice','Next Speaker: Alice\nNext Speaker: Bob','Next Speaker : Alice','Next　Speaker: 中文','','  ',':','a:b:c','The next speaker is: Alice\n']
JS=['{"Storyline Quality":{"flaws":[{"instance":"x","type":"t","severity":3}]}}','```json\n{"Storyline Quality":{"flaws":[]}}\n```','{"Storyline Quality":{"flaws":[{"severity":null},{"severity":2},{"severity":"3"},{"severity":1.5},{"severity":true},{}]}}','text before {"Storyline Quality":{"flaws":[{"severity":5}]}} after {"x":1}','{"Storyline Quality":{"flaws":[{"instance":"multi\nline","severity":2}]}}','{"Storyline Quality":{"flaws":[{"severity":2},]}}','{"Storyline Quality": {"flaws": "abc"}}','{"Storyline Quality": {"flaws": ""}}','{"Storyline Quality": {"flaws": {}}}','{"Storyline Quality": {"flaws": {"a":1}}}','{"Storyline Quality": ["flaws"]}','{"Storyline Quality": "flaws"}','{"Anthropomorphism":{"flaws":[]}}','{"Other":{"flaws":[]}}','{}','[]','null','0','"str"','[1,2,3]','{"Storyline Quality":{"flaws":[{"severity":1e2},{"severity":100000000000000000000}]}}','{"Storyline Quality":{"flaws":[{"severity":-3}]}}','{"Storyline Quality":{"flaws":[{"severity":NaN},{"severity":Infinity}]}}','{"Storyline Quality":{"flaws":[{"instance":"unescaped "quote" here","severity":2}]}}','{"a":{"b":1}} {"Storyline Quality":{"flaws":[{"severity":2}]}}','{"Storyline Quality":{"flaws":[{"severity":2}]},"Storyline Quality":{"flaws":[{"severity":4}]}}','no json here','','"a\nb"','{"Storyline Quality":{"flaws":null}}','{"Storyline Quality":{"flaws":[1]}}','{"Storyline Quality":{"flaws":[{"severity":3,"severity":null}]}}','1 2 3 {"Storyline Quality":{"flaws":[]}}','{"Storyline Quality":{"flaws":[{"severity":0}]}}','{"Storyline Quality":{"flaws":[]}}','{"Storyline Quality":{"flaws":[{"instance":"\\ud83d\\ude00","severity":1}]}}']
out={'rit':[],'add':[],'nsp':[],'json':[],'own':[],'fix':[]}
for i in range(6000):
    s=''.join(rnd.choice(TH) for _ in range(rnd.randint(0,8)))
    out['rit'].append([s,remove_inner_thoughts(s)])
    out['add'].append([s,'Alice',add_speaker_name(s,'Alice')])
    o=own(s,'Alice'); out['own'].append([s,o])
for s in NS: out['nsp'].append([s,parse_nsp(s)])
for i in range(1500):
    s=''.join(rnd.choice(NS+['\n',' ','x ']) for _ in range(rnd.randint(1,4))); out['nsp'].append([s,parse_nsp(s)])
for s in JS:
    for R in (0,1,7,13,20): out['json'].append([s,R,read(s,'Storyline Quality',R)])
W=['the','quick','brown','fox','jumps','over','lazy','dog','(nods)','!','?',',','.','中','文','あ','ア','1','23','A'*120,'B'*100,'-','"',"it's"]
for i in range(1500):
    base=[rnd.choice(W) for _ in range(rnd.randint(0,12))]
    rep=rnd.choice([1,1,2,5,6]); tail=[rnd.choice(W) for _ in range(rnd.randint(0,5))]
    if rnd.random()<.3: seq=base+tail+base*rep
    else: seq=base*rep+tail
    s=' '.join(seq); out['fix'].append([s,fix(s)])
json.dump(out,open(sys.argv[1],'w'),ensure_ascii=False)
print({k:len(v) for k,v in out.items()})
print('NSP parse of typical replies:')
for s in NS[:24]:
    print(' ',json.dumps(s),'->',json.dumps(parse_nsp(s)))