#!/usr/bin/env python3
"""Re-run the measurement behind Phayan's what-ka-means-in-thai post.

Corpus: every Thai headword in kaikki.org's Wiktionary extract that carries a
Royal Institute (RTGS) romanisation and an IPA transcription. Spellings are
grouped by romanisation, once under RTGS, which writes no tone and no vowel
length, and once under Paiboon, which writes both. A word is ambiguous when
another spelling lands on the same string. "Common words" are the words of
Phayan's own word list that carry a Thai National Corpus rank; the list is
stored in dataset.json so the script needs nothing but the public dump.

Pages: re-fetch every page the post cites and check its quoted sentence is
still there.

  check.py               download the dump, re-measure, compare with dataset.json
                         and dataset.csv, then check every cited page
  check.py --recheck     the same, but the dump download is skipped if it takes
                         longer than 60 s
  check.py --dump PATH   use a local copy of the dump instead of downloading it
  check.py --write       regenerate the results in dataset.json and dataset.csv

Exit 0 when everything holds, 1 otherwise. Standard library only, plus
pdftotext for the one PDF page.
"""

import collections
import csv
import hashlib
import html
import json
import os
import re
import subprocess
import sys
import tempfile
import time
import urllib.request

HERE = os.path.dirname(os.path.abspath(__file__))
DATASET_JSON = os.path.join(HERE, 'dataset.json')
DATASET_CSV = os.path.join(HERE, 'dataset.csv')

DUMP_URL = 'https://kaikki.org/dictionary/Thai/kaikki.org-dictionary-Thai.jsonl'
POST_WORDS = ['กา', 'ขา', 'ค่า', 'ค่ะ', 'คะ', 'ครับ', 'คับ']
RECHECK_BUDGET_S = 60
# When Wiktionary has changed since the recorded dump, corpus shares may drift
# by a few words without the post's claims changing.
DRIFT_PP = 1.0

TONES = {'˧': 'mid', '˨˩': 'low', '˥˩': 'falling', '˦˩': 'falling',
         '˦˥': 'high', '˩˩˦': 'rising', '˨˩˦': 'rising'}
TONE_LETTERS = set('˥˦˧˨˩')

CSV_FIELDS = ['word', 'in_post', 'rtgs', 'paiboon', 'ipa', 'tone', 'glosses',
              'rtgs_group_size', 'rtgs_group', 'paiboon_group_size', 'tnc_rank']


def tone_runs(ipa):
    runs, cur = [], ''
    for c in ipa:
        if c in TONE_LETTERS:
            cur += c
        elif cur:
            runs.append(cur)
            cur = ''
    if cur:
        runs.append(cur)
    return runs


def fetch(url, dest, deadline=None):
    req = urllib.request.Request(url, headers={'User-Agent': 'Mozilla/5.0'})
    digest, size, start = hashlib.sha256(), 0, time.time()
    with urllib.request.urlopen(req, timeout=120) as resp, open(dest, 'wb') as fh:
        while True:
            if deadline and time.time() - start > deadline:
                raise TimeoutError(f'download passed {deadline} s')
            chunk = resp.read(1 << 20)
            if not chunk:
                break
            digest.update(chunk)
            fh.write(chunk)
            size += len(chunk)
    return digest.hexdigest(), size


def sha256_file(path):
    digest = hashlib.sha256()
    with open(path, 'rb') as fh:
        for chunk in iter(lambda: fh.read(1 << 20), b''):
            digest.update(chunk)
    return digest.hexdigest()


def parse_dump(path):
    """word -> romanisations, IPA and one gloss per entry; the first entry with both wins."""
    words = {}
    glosses = collections.defaultdict(list)
    with open(path, encoding='utf-8') as fh:
        for line in fh:
            if not line.strip():
                continue
            e = json.loads(line)
            word = e.get('word')
            if not word or re.search(r'\s', word):
                continue
            ri = pb = ipa = None
            for s in e.get('sounds') or []:
                tag = s.get('raw_tags')
                if ri is None and tag == ['Royal Institute'] and s.get('roman'):
                    ri = s['roman']
                if pb is None and tag == ['Paiboon'] and s.get('roman'):
                    pb = s['roman']
                if ipa is None and s.get('ipa'):
                    ipa = s['ipa']
            for sense in e.get('senses') or []:
                g = next((g for g in sense.get('glosses') or [] if isinstance(g, str) and g), '')
                if g:
                    if g not in glosses[word]:
                        glosses[word].append(g)
                    break
            if ri and pb and ipa and word not in words:
                runs = tone_runs(ipa)
                words[word] = {'rtgs': ri, 'paiboon': pb, 'ipa': ipa,
                               'syllables': len(runs),
                               'tone': TONES.get(runs[0], '') if len(runs) == 1 else ''}
    for word, rec in words.items():
        rec['glosses'] = ' | '.join(glosses[word])
    return words


def groups(words, key):
    out = collections.defaultdict(list)
    for w, rec in words.items():
        out[rec[key]].append(w)
    return {k: sorted(v) for k, v in out.items()}


def stats(words, key):
    g = groups(words, key)
    multi = [v for v in g.values() if len(v) > 1]
    amb = sum(len(v) for v in multi)
    return {'words': len(words), 'distinctRomanisations': len(g),
            'groupsWithTwoOrMoreWords': len(multi), 'wordsInSuchAGroup': amb,
            'percentAmbiguous': round(100 * amb / len(words), 1) if words else 0.0}


def measure(words, ranks):
    common = {w: r for w, r in words.items() if w in ranks}
    mono = {w: r for w, r in words.items() if r['syllables'] == 1}
    rg, pg = groups(common, 'rtgs'), groups(common, 'paiboon')
    rows, seen = [], set()
    for pw in POST_WORDS:
        for w in rg[common[pw]['rtgs']]:
            if w in seen:
                continue
            seen.add(w)
            r = common[w]
            rows.append({
                'word': w, 'in_post': 'yes' if w in POST_WORDS else 'no',
                'rtgs': r['rtgs'], 'paiboon': r['paiboon'], 'ipa': r['ipa'],
                'tone': r['tone'], 'glosses': r['glosses'],
                'rtgs_group_size': len(rg[r['rtgs']]),
                'rtgs_group': ' '.join(rg[r['rtgs']]),
                'paiboon_group_size': len(pg[r['paiboon']]),
                'tnc_rank': ranks[w],
            })
    top = sorted(rg.items(), key=lambda kv: (-len(kv[1]), kv[0]))[:5]
    return {
        'corpus': {
            'rtgs': {'allWords': stats(words, 'rtgs'), 'monosyllabic': stats(mono, 'rtgs'),
                     'commonWords': stats(common, 'rtgs')},
            'paiboon': {'allWords': stats(words, 'paiboon'), 'monosyllabic': stats(mono, 'paiboon'),
                        'commonWords': stats(common, 'paiboon')},
        },
        'largestRtgsGroupsCommonWords': [{'romanisation': k, 'words': v} for k, v in top],
        'rows': rows,
    }


# Page text loses tags, so "ครับ (khrap)" arrives as "ครับ ( khrap )"; both
# sides are normalised the same way before comparison.
QUOTES = str.maketrans({'‘': "'", '’': "'", '“': '"', '”': '"', '–': '-',
                        '—': '-', ' ': ' ', ' ': ' '})


def norm(text):
    t = ' '.join(html.unescape(str(text)).translate(QUOTES).split())
    t = re.sub(r'\s+([),.;:])', r'\1', t)
    return re.sub(r'\(\s+', '(', t)


def page_text(url):
    req = urllib.request.Request(url, headers={'User-Agent': 'Mozilla/5.0'})
    with urllib.request.urlopen(req, timeout=90) as resp:
        raw = resp.read()
    if url.lower().endswith('.pdf'):
        with tempfile.NamedTemporaryFile(suffix='.pdf', delete=False) as fh:
            fh.write(raw)
        try:
            return subprocess.run(['pdftotext', '-layout', fh.name, '-'],
                                  capture_output=True, check=False).stdout.decode('utf-8', 'replace')
        finally:
            os.unlink(fh.name)
    text = raw.decode('utf-8', 'replace')
    text = re.sub(r'<script.*?</script>|<style.*?</style>', ' ', text, flags=re.S)
    text = re.sub(r'<sup[^>]*>.*?</sup>', ' ', text, flags=re.S)
    return re.sub(r'<[^>]+>', ' ', text)


def check_pages(pages, failures):
    cache = {}
    for p in pages:
        try:
            if p['url'] not in cache:
                cache[p['url']] = norm(page_text(p['url']))
        except Exception as exc:  # noqa: BLE001 - report, never fake a pass
            print(f"  {p['id']}: UNREACHABLE {type(exc).__name__}")
            failures.append(f"{p['id']}: {p['url']} unreachable: {exc}")
            continue
        found = norm(p['quote']) in cache[p['url']]
        print(f"  {p['id']}: {'found' if found else 'MISSING'}")
        if not found:
            failures.append(f"{p['id']}: quote no longer on {p['url']}")


def compare(path, expected, actual, failures, drift):
    if isinstance(expected, dict):
        for k in expected:
            if k not in actual:
                failures.append(f'{path}.{k}: missing from the measurement')
            else:
                compare(f'{path}.{k}', expected[k], actual[k], failures, drift)
    elif isinstance(expected, list):
        if len(expected) != len(actual):
            failures.append(f'{path}: stored {len(expected)} items, measured {len(actual)}')
            return
        for i, (e, a) in enumerate(zip(expected, actual)):
            compare(f'{path}[{i}]', e, a, failures, drift)
    elif drift and path.startswith('results.corpus') and path.endswith('percentAmbiguous'):
        if abs(expected - actual) > DRIFT_PP:
            failures.append(f'{path}: stored {expected}, measured {actual}')
        elif expected != actual:
            print(f'  drift within {DRIFT_PP} pp: {path} {expected} -> {actual}')
    elif drift and path.startswith('results.corpus'):
        if expected != actual:
            print(f'  drift: {path} {expected} -> {actual}')
    elif norm(expected) != norm(actual):
        failures.append(f'{path}: stored {expected!r}, measured {actual!r}')


def read_csv():
    with open(DATASET_CSV, encoding='utf-8') as fh:
        return list(csv.DictReader(fh))


def write_csv(rows):
    with open(DATASET_CSV, 'w', encoding='utf-8', newline='') as fh:
        w = csv.DictWriter(fh, fieldnames=CSV_FIELDS)
        w.writeheader()
        w.writerows(rows)


def main():
    args = sys.argv[1:]
    recheck, write = '--recheck' in args, '--write' in args
    dump = args[args.index('--dump') + 1] if '--dump' in args else None
    dataset = json.load(open(DATASET_JSON, encoding='utf-8'))
    source = next(i for i in dataset['inputs'] if i['id'] == 'kaikki-thai')
    ranks = next(i for i in dataset['inputs'] if i['id'] == 'phayan-word-list')['words']
    failures = []

    with tempfile.TemporaryDirectory() as tmp:
        measured = None
        try:
            if dump:
                sha = sha256_file(dump)
                print(f'dump {dump} sha256 {sha}')
            else:
                dump = os.path.join(tmp, 'thai.jsonl')
                print(f'downloading {DUMP_URL}')
                sha, size = fetch(DUMP_URL, dump, RECHECK_BUDGET_S if recheck else None)
                print(f'  sha256 {sha} ({size} bytes)')
            measured = measure(parse_dump(dump), ranks)
        except TimeoutError as exc:
            print(f'lab step skipped: {exc}')
        except Exception as exc:  # noqa: BLE001 - report, never fake a pass
            print(f'SOURCE UNREACHABLE: {type(exc).__name__}: {exc}')
            failures.append(f'dump: {exc}')

    if measured and write:
        source['sha256'] = sha
        dataset['results'] = measured
        with open(DATASET_JSON, 'w', encoding='utf-8') as fh:
            json.dump(dataset, fh, ensure_ascii=False, indent=2)
            fh.write('\n')
        write_csv(measured['rows'])
        print('wrote dataset.json and dataset.csv')
    elif measured:
        drift = sha != source['sha256']
        if drift:
            print(f'note: the dump changed since {dataset["checkedAt"]} '
                  f'(recorded sha256 {source["sha256"]}); corpus totals may drift')
        compare('results', dataset['results'], measured, failures, drift)

    stored_rows = [{k: str(v) for k, v in r.items()} for r in dataset['results']['rows']]
    if read_csv() != stored_rows:
        failures.append('dataset.csv does not match results.rows in dataset.json')

    print('cited pages')
    check_pages(dataset['pages'], failures)

    if failures:
        print(f'\nFAIL ({len(failures)})')
        for f in failures:
            print(f'  {f}')
        return 1

    res = measured or dataset['results']
    for scheme in ('rtgs', 'paiboon'):
        for subset, s in res['corpus'][scheme].items():
            print(f"{scheme:<8} {subset:<13} {s['words']:>6} words -> "
                  f"{s['distinctRomanisations']:>6} strings, {s['percentAmbiguous']}% ambiguous")
    for r in res['rows']:
        if r['in_post'] == 'yes':
            print(f"  {r['word']}  rtgs {r['rtgs']:<6} (shared by {r['rtgs_group_size']})  "
                  f"paiboon {r['paiboon']:<6} (shared by {r['paiboon_group_size']})  {r['tone']}")
    print('\nOK: every stored value and every quote holds')
    return 0


if __name__ == '__main__':
    sys.exit(main())