#!/usr/bin/env python3
"""Summarize local browser telemetry and compare small logits; stdlib only.

Never downloads, converts, or loads model weights. Each reference vector is
166,144 float32 values (about 649 KiB). Missing evidence stays unavailable.
"""
from __future__ import annotations

import argparse
from array import array
from collections import Counter, defaultdict
from datetime import datetime, timezone
import hashlib
import json
import math
from pathlib import Path
import statistics
import sys
from urllib.parse import urlparse, unquote

ROOT = Path(__file__).resolve().parents[1]
CUSTOM_REPO = 'borkiss/Nanbeige4.2-3B-WebGPU'
VOCAB = 166144
METRICS = ('ttftMs', 'elapsedMs', 'prefillTokensPerSecond', 'decodeTokensPerSecond')


def canonical(value):
    return json.dumps(value, sort_keys=True, ensure_ascii=False, separators=(',', ':'))


def digest(value):
    return hashlib.sha256(canonical(value).encode()).hexdigest()


def finite(value):
    return isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(value)


def stats(values):
    values = [float(x) for x in values if finite(x)]
    if not values:
        return {'count': 0, 'median': None, 'min': None, 'max': None}
    return {'count': len(values), 'median': statistics.median(values), 'min': min(values), 'max': max(values)}


def relative(path, root):
    try:
        return str(path.resolve().relative_to(root.resolve()))
    except ValueError:
        return str(path.resolve())


def read_json(path):
    with path.open() as stream:
        return json.load(stream)


def artifact_variant(label):
    if not isinstance(label, str) or not label:
        return None
    # Packed is a byte-layout change, not a different quantization algorithm.
    label = label.removesuffix('-packed')
    return label if label.startswith('webgpu-') else 'webgpu-' + label


def identify_variant(run):
    if run.get('engine') != 'custom-wgsl' or run.get('model') != CUSTOM_REPO:
        return None, 'No matching quantized reference registered for this engine/model.'
    explicit = artifact_variant(run.get('variant'))
    from_url = None
    base = run.get('baseUrl') or run.get('manifestUrl')
    if base:
        parsed = urlparse(base)
        parts = [unquote(x) for x in parsed.path.split('/') if x]
        if parsed.hostname != 'huggingface.co' or len(parts) < 5 or '/'.join(parts[:2]) != CUSTOM_REPO or parts[2] != 'resolve':
            return None, 'Artifact URL is not an expected pinned custom-model URL.'
        if run.get('revision') and parts[3] != run['revision']:
            return None, 'Reported revision disagrees with artifact URL.'
        from_url = artifact_variant(parts[4])
    if explicit and from_url and explicit != from_url:
        return None, 'Reported variant disagrees with artifact URL.'
    value = explicit or from_url
    if not value:
        return None, 'No exact artifact variant is recorded; runKind/prompt are not used to guess it.'
    return value, None


def load_reference_index(folder, root):
    entries, notes = {}, []
    # Prefer the top-level copy, retaining the older nested download layout only as fallback.
    files = sorted(folder.glob('*-quality.json')) + sorted((folder / 'results').glob('*-quality.json'))
    for path in files:
        try:
            data = read_json(path)
            variant = data.get('variant')
            if not variant:
                notes.append({'file': relative(path, root), 'error': 'Missing exact variant label.'})
                continue
            if variant in entries:
                previous = entries[variant]['quality']
                if canonical(previous) != canonical(data):
                    notes.append({'file': relative(path, root), 'error': 'Conflicting duplicate quality JSON; primary copy retained.'})
                continue
            ids = defaultdict(list)
            for index, generation in enumerate(data.get('generations', [])):
                token_ids = generation.get('inputIds')
                if isinstance(token_ids, list):
                    ids[canonical(token_ids)].append(index)
            entries[variant] = {'path': path, 'quality': data, 'inputIndex': ids}
        except (OSError, ValueError, TypeError) as error:
            notes.append({'file': relative(path, root), 'error': str(error)})
    return entries, notes


def load_logits(path):
    if path.stat().st_size != VOCAB * 4:
        raise ValueError(f'Expected {VOCAB * 4} bytes, found {path.stat().st_size}.')
    values = array('f')
    with path.open('rb') as stream:
        values.fromfile(stream, VOCAB)
    if sys.byteorder != 'little':
        values.byteswap()
    if not all(math.isfinite(x) for x in values):
        raise ValueError('Reference contains non-finite logits.')
    return values


def compare_logits(actual, reference):
    if len(actual) != VOCAB or len(reference) != VOCAB:
        raise ValueError(f'Logit length mismatch; expected {VOCAB}.')
    if not all(finite(x) for x in actual):
        raise ValueError('Browser logits contain non-finite or non-numeric values.')
    dot = math.fsum(a * b for a, b in zip(actual, reference))
    aa = math.fsum(a * a for a in actual)
    bb = math.fsum(b * b for b in reference)
    diffs = [abs(a - b) for a, b in zip(actual, reference)]
    ai = max(range(VOCAB), key=actual.__getitem__)
    bi = max(range(VOCAB), key=reference.__getitem__)
    amax, bmax = actual[ai], reference[bi]
    alse = amax + math.log(math.fsum(math.exp(x - amax) for x in actual))
    blse = bmax + math.log(math.fsum(math.exp(x - bmax) for x in reference))
    kl = math.fsum(math.exp(b - blse) * ((b - blse) - (a - alse)) for a, b in zip(actual, reference))
    cosine = dot / math.sqrt(aa * bb) if aa and bb else None
    return {
        'logitCount': VOCAB, 'maxAbs': max(diffs), 'meanAbs': statistics.fmean(diffs),
        'rmse': math.sqrt(math.fsum(x * x for x in diffs) / VOCAB),
        'cosine': cosine, 'klReferenceToBrowser': max(0.0, kl),
        'top1Equal': ai == bi, 'browserTop1': ai, 'referenceTop1': bi,
    }


def compare_token_prefix(actual, reference):
    actual = actual if isinstance(actual, list) else []
    reference = reference if isinstance(reference, list) else []
    length = min(len(actual), len(reference))
    mismatches = [i for i in range(length) if actual[i] != reference[i]]
    return {'comparedTokens': length, 'matchingTokens': length - len(mismatches),
            'equalAtCommonLength': bool(length) and not mismatches,
            'browserLength': len(actual), 'referenceLength': len(reference),
            'firstDifferentIndex': mismatches[0] if mismatches else None,
            'browserFullyCovered': bool(actual) and len(actual) <= len(reference)}


def match_reference(run, index, folder, source_revision):
    variant, error = identify_variant(run)
    if error:
        return {'status': 'unavailable', 'reason': error}, None
    entry = index.get(variant)
    if entry is None:
        return {'status': 'unavailable', 'variant': variant, 'reason': 'Matching variant quality JSON has not been downloaded.'}, None
    quality = entry['quality']
    if run.get('sourceRevision') and source_revision and run['sourceRevision'] != source_revision:
        return {'status': 'unavailable', 'variant': variant, 'reason': 'Source revision disagrees with reference-depths provenance.'}, None
    if run.get('weightBytes') is not None and quality.get('weightBytes') is not None and run['weightBytes'] != quality['weightBytes']:
        return {'status': 'unavailable', 'variant': variant, 'reason': 'Weight byte count differs from matching quality artifact.'}, None
    indices = entry['inputIndex'].get(canonical(run.get('inputIds')), [])
    if len(indices) != 1:
        return {'status': 'unavailable', 'variant': variant, 'reason': 'Input token IDs have no unique exact match in variant quality JSON.'}, None
    case = indices[0]
    name = f'{variant}-logits-{case}.f32'
    candidates = [entry['path'].parent / name, folder / name, folder / 'results' / name]
    path = next((x for x in candidates if x.is_file()), None)
    if path is None:
        return {'status': 'unavailable', 'variant': variant, 'promptIndex': case, 'reason': 'Matching reference logit vector has not been downloaded.', 'expectedFile': name}, None
    return {
        'status': 'matched', 'variant': variant, 'promptIndex': case,
        'qualityFile': str(entry['path']), 'referenceFile': str(path),
        'matchBasis': 'Exact model/variant identity, compatible recorded source revision and byte count, exact input token IDs in quality generation order.',
        'artifactHashVerified': False,
    }, path


def run_metadata(run, line, offset, event):
    vector = run.get('firstLogits')
    capture = run.get('captureLogits', bool(vector))
    device = run.get('device') or {}
    profile = {
        'context': run.get('context'), 'chunkSize': run.get('chunkSize'),
        'f16': device.get('f16'), 'subgroups': device.get('subgroups'),
        'subgroupSize': device.get('subgroupSize'), 'packed': run.get('packed'),
        'tuning': run.get('tuning'), 'effectiveProfile': run.get('effectiveProfile'),
        'captureLogits': capture, 'ignoreEos': run.get('ignoreEos'), 'thinking': run.get('thinking'),
    }
    kind = run.get('runKind') or 'unrecorded'
    warmup = bool(run.get('warmup')) or 'warmup' in kind.lower()
    return {
        'runId': f'line-{line}', 'eventLine': line, 'eventByteOffset': offset,
        'timestamp': run.get('timestamp'), 'received': event.get('received'),
        'engine': run.get('engine'), 'build': run.get('build') or 'unrecorded',
        'runtime': run.get('runtime'), 'runKind': kind,
        'model': run.get('model'), 'revision': run.get('revision'),
        'sourceRevision': run.get('sourceRevision'), 'variant': run.get('variant'),
        'baseUrl': run.get('baseUrl'), 'device': device, 'userAgent': run.get('userAgent'),
        'profile': profile, 'prompt': run.get('prompt'), 'inputIdsSha256': digest(run.get('inputIds')),
        'promptTokens': run.get('promptTokens'), 'generatedTokens': run.get('generatedTokens'),
        'outputIdsSha256': digest(run.get('outputIds')), 'cancelled': bool(run.get('cancelled')),
        'explicitWarmup': warmup, 'repeatIndex': run.get('repeatIndex'),
        'loadMs': run.get('loadMs'), 'weightBytes': run.get('weightBytes'), 'kvBytes': run.get('kvBytes'),
        'firstLogitsCount': len(vector) if isinstance(vector, list) else 0,
        **{metric: run.get(metric) for metric in METRICS},
    }


def summarize_references(index, folder, root):
    output = {'quantization': [], 'depths': None, 'corpusScreenings': []}
    for variant, entry in sorted(index.items()):
        q = entry['quality']
        generations = q.get('generations', [])
        output['quantization'].append({
            'variant': variant, 'file': relative(entry['path'], root), 'weightBytes': q.get('weightBytes'),
            'perplexity': q.get('perplexity'), 'scoredTokens': sum(x.get('tokens', 0) for x in q.get('nll', [])),
            'generationCount': len(generations),
            'sourceTop1Matches': sum(x.get('top1Equal') is True for x in generations),
            'logitCosineVsSource': stats(x.get('logitCosine') for x in generations),
            'klVsSource': stats(x.get('klDivergence') for x in generations),
        })
    depths = folder / 'reference-depths.json'
    if depths.is_file():
        d = read_json(depths)
        output['depths'] = {'file': relative(depths, root), 'sourceRevision': d.get('sourceRevision'), 'dtype': d.get('dtype'), 'gpu': d.get('gpu'), 'depths': [{key: x.get(key) for key in ('loops', 'perplexity', 'cacheParity')} for x in d.get('depths', [])]}
    for path in sorted(folder.glob('wikitext-screening*.json')):
        data = read_json(path)
        output['corpusScreenings'].append({'file': relative(path, root), **data})
        output['corpusScreenings'][-1].pop('tokenIds', None)
    return output


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--events', type=Path, default=ROOT / 'results/browser-events.jsonl')
    parser.add_argument('--references', type=Path, default=ROOT / 'results/reference')
    parser.add_argument('--output-dir', type=Path, default=ROOT / 'results')
    args = parser.parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=True)
    index, reference_notes = load_reference_index(args.references, ROOT)
    ref_summary = summarize_references(index, args.references, ROOT)
    source_revision = (ref_summary.get('depths') or {}).get('sourceRevision')
    runs, parity, cache_parity, errors, malformed, kernel_tests = [], [], [], [], [], []
    counts, groups = Counter(), defaultdict(list)
    previous_vectors = {}  # One small vector per exact config and tokenized prompt, for chunk comparisons.
    event_hash = hashlib.sha256()
    snapshot_size = args.events.stat().st_size if args.events.is_file() else 0
    consumed = 0
    if args.events.is_file():
        with args.events.open('rb') as stream:
            for line_number, raw in enumerate(stream, 1):
                # Ignore later appends, so hashes and summaries describe a reproducible prefix.
                if consumed + len(raw) > snapshot_size:
                    break
                offset = consumed
                consumed += len(raw)
                event_hash.update(raw)
                try:
                    event = json.loads(raw)
                except (ValueError, UnicodeDecodeError) as error:
                    malformed.append({'line': line_number, 'byteOffset': offset, 'error': str(error), 'incompleteTail': not raw.endswith(b'\n')})
                    continue
                counts[event.get('type', 'unknown')] += 1
                if event.get('type') in ('kernel-tests', 'gpu-candidate-tests'):
                    kernel_tests.append({'eventLine': line_number, **event})
                if event.get('type') == 'error':
                    errors.append({key: event.get(key) for key in ('received', 'message', 'cancelled', 'requestRejected')})
                if event.get('type') != 'result' or not isinstance(event.get('result'), dict):
                    continue
                run = event['result']
                meta = run_metadata(run, line_number, offset, event)
                runs.append(meta)
                group_key = {key: meta[key] for key in ('engine', 'build', 'runtime', 'runKind', 'model', 'revision', 'variant', 'device', 'userAgent', 'profile', 'inputIdsSha256', 'promptTokens', 'generatedTokens', 'explicitWarmup')}
                group_id = digest(group_key)[:16]
                meta['groupId'] = group_id
                groups[group_id].append(meta)
                vector = run.get('firstLogits')
                if not isinstance(vector, list) or not vector:
                    continue
                match, reference = match_reference(run, index, args.references, source_revision)
                record = {'runId': meta['runId'], 'build': meta['build'], 'runKind': meta['runKind'], 'revision': meta['revision'], 'profile': meta['profile'], **match}
                for field in ('qualityFile', 'referenceFile'):
                    if field in record:
                        record[field] = relative(Path(record[field]), ROOT)
                if meta['cancelled']:
                    record.update(status='unavailable', reason='Cancelled generation excluded from parity evidence.')
                elif reference:
                    try:
                        record['referenceSha256'] = hashlib.sha256(reference.read_bytes()).hexdigest()
                        record['metrics'] = compare_logits(vector, load_logits(reference))
                        reference_generation = index[record['variant']]['quality']['generations'][record['promptIndex']]
                        record['generatedPrefix'] = compare_token_prefix(run.get('outputIds'), reference_generation.get('outputIds'))
                        record['status'] = 'compared'
                    except (OSError, ValueError, OverflowError) as error:
                        record.update(status='unavailable', reason=str(error))
                parity.append(record)
                # Compare distinct prefill chunking of the exact same variant, build, device and input.
                if record['status'] == 'compared':
                    comparable = dict(group_key)
                    comparable.pop('runKind');comparable.pop('generatedTokens');comparable.pop('explicitWarmup')
                    comparable['profile'] = {k: v for k, v in meta['profile'].items() if k not in ('chunkSize', 'captureLogits', 'ignoreEos')}
                    key = canonical(comparable)
                    prior = previous_vectors.get(key)
                    if prior and prior['chunkSize'] != run.get('chunkSize'):
                        cache_parity.append({'runId': meta['runId'], 'referenceRunId': prior['runId'], 'chunkSize': run.get('chunkSize'), 'referenceChunkSize': prior['chunkSize'], 'variant': record['variant'], 'metrics': compare_logits(vector, prior['values']), 'generatedPrefix': compare_token_prefix(run.get('outputIds'), prior.get('outputIds')), 'scope': 'Same input, different prefill chunking; browser float32-logit comparison.'})
                    previous_vectors[key] = {'runId': meta['runId'], 'chunkSize': run.get('chunkSize'), 'values': array('f', vector), 'outputIds': run.get('outputIds')}
    summaries = []
    for group_id, rows in groups.items():
        sample = rows[0]
        complete = [r for r in rows if not r['cancelled']]
        summaries.append({
            'groupId': group_id, **{k: sample[k] for k in ('engine', 'build', 'runtime', 'runKind', 'model', 'revision', 'variant', 'device', 'userAgent', 'profile', 'prompt', 'inputIdsSha256', 'promptTokens', 'generatedTokens', 'explicitWarmup')},
            'runs': len(rows), 'completedRuns': len(complete), 'cancelledRuns': len(rows) - len(complete),
            'runIds': [r['runId'] for r in rows],
            'performanceCandidate': bool(complete) and not sample['explicitWarmup'] and not sample['profile']['captureLogits'] and 'validation' not in sample['runKind'],
            'metrics': {metric: stats(r[metric] for r in complete) for metric in METRICS},
        })
    now = datetime.now(timezone.utc).isoformat()
    source = {'path': relative(args.events, ROOT), 'snapshotBytes': snapshot_size, 'consumedBytes': consumed, 'consumedPrefixSha256': event_hash.hexdigest(), 'eventCounts': dict(counts)}
    summary = {
        'schemaVersion': 1, 'generatedAt': now, 'source': source,
        'method': 'Medians/min/max across completed runs sharing build, runKind, recorded device/profile, exact prompt IDs and generated-token count. Explicit warmups and captured-logit diagnostics remain separate. No implicit warmup removal or hardware identity inference.',
        'limitations': ['Legacy missing build/profile/runKind fields remain unrecorded/null.', 'loadMs is repeated session initialization metadata, not a separate fresh load measurement per generation.', 'Timings are worker wall time; no GPU-only timing or physical memory measurement is inferred.', 'A performanceCandidate is a filter, not a claim that runs were warmed, foreground, thermally stable or sufficient in number.'],
        'runs': runs, 'groups': summaries, 'errors': errors, 'malformedEvents': malformed, 'kernelTests': kernel_tests,
        'referenceNotes': reference_notes, 'referenceSummaries': ref_summary,
    }
    parity_output = {
        'schemaVersion': 1, 'generatedAt': now, 'source': source,
        'method': 'All 166,144 first-token logits compared to the exactly named quantization variant and exact token-ID case. KL direction is reference to browser. No threshold or pass claim is invented.',
        'limitations': ['Quality JSONs may not carry quantized artifact hashes; variant/source/input matching is metadata-based, not cryptographic weight-identity verification.', 'FP32 dequantized remote reference differs numerically from browser FP16 KV and GPU reductions.', 'Missing references, invalid vectors, and cancelled generations are explicitly unavailable; ORT is never substituted with a custom quant reference.'],
        'capturedRuns': len(parity), 'comparedRuns': sum(r['status'] == 'compared' for r in parity),
        'unavailableRuns': sum(r['status'] != 'compared' for r in parity),
        'runs': parity, 'prefillChunkComparisons': cache_parity,
    }
    for name, payload in [('browser-summary.json', summary), ('browser-parity.json', parity_output)]:
        path = args.output_dir / name
        temporary = path.with_suffix('.tmp')
        temporary.write_text(json.dumps(payload, ensure_ascii=False, indent=2, allow_nan=False) + '\n')
        temporary.replace(path)
    print(json.dumps({'runs': len(runs), 'groups': len(summaries), 'comparedLogitRuns': parity_output['comparedRuns'], 'unavailableLogitRuns': parity_output['unavailableRuns'], 'prefillChunkComparisons': len(cache_parity), 'outputDirectory': str(args.output_dir)}, indent=2))


if __name__ == '__main__':
    main()