#!/usr/bin/env python3
"""Correlate content-free Gateway/ESP32 markers, never speech or Agent text.
Clock alignment is relative to the first common sequence, NOT one-way latency.
Normal omitted serial samples are unknown, never inferred as missing chunks.
"""
import argparse
import csv
import io
import json
import re
from pathlib import Path

ANSI = re.compile(r'\x1b\[[0-?]*[ -/]*[@-~]')
KV = re.compile(r'\b([a-z_]+)=([^\s]+)')


def fields(line):
    result = {}
    for key, value in KV.findall(line):
        try:
            result[key] = float(value) if '.' in value else int(value)
        except ValueError:
            result[key] = value
    return result


def correlate(gateway_text, serial_text, speech_id):
    gateway = []
    malformed = 0
    for line in gateway_text.splitlines():
        if '[SPEECH_TRACE] ' not in line:
            continue
        try:
            record = json.loads(line.split('[SPEECH_TRACE] ', 1)[1])
        except (ValueError, TypeError):
            malformed += 1
            continue
        if record.get('speech_id') == speech_id:
            gateway.append(record)
    sends = [r for r in gateway if r.get('event') == 'send']
    gs = next((r for r in reversed(gateway) if r.get('event') == 'summary'), {})
    chunks = {r['sequence']: r for r in sends if r.get('message_type') == 'speech_chunk'}
    gend = next((r for r in reversed(sends) if r.get('message_type') == 'speech_end'), {})
    expected = gend.get('chunks', gs.get('completed_chunks') if gs.get('speech_end_sent') else None)
    expected_bytes = gs.get('raw_pcm_bytes')
    lines = [ANSI.sub('', line) for line in serial_text.splitlines()]
    # Bound to the selected stream so an earlier boot/disconnect isn't attributed to it.
    start = next((i for i, l in enumerate(lines) if '[SPEECH_RX_EVENT]' in l
                  and fields(l).get('event') == 'start' and fields(l).get('id') == speech_id), None)
    stop = len(lines)
    if start is not None:
        stop = next((i for i in range(start+1, len(lines)) if '[SPEECH_RX_EVENT]' in lines[i]
                     and fields(lines[i]).get('event') == 'start'), len(lines))
    scoped = lines[start:stop] if start is not None else []
    rx, summaries, events = {}, [], []
    generation = None
    valid_end = False
    for line in scoped:
        if not any(marker in line for marker in ('[SPEECH_RX]', '[SPEECH_RX_SUMMARY]', '[SPEECH_RX_EVENT]', '[SPEECH] start', '[SPEECH] end')):
            continue
        f = fields(line)
        if f.get('id') != speech_id:
            continue
        if '[SPEECH_RX]' in line:
            # 'last' repeats an earlier completion: preserve the original observation.
            rx.setdefault(f['seq'], f)
        if '[SPEECH_RX_SUMMARY]' in line:
            summaries.append(f)
        if '[SPEECH_RX_EVENT]' in line:
            events.append(f)
        if '[SPEECH] start' in line:
            generation = f.get('generation')
        if '[SPEECH] end' in line:
            valid_end = f.get('valid') == 1 and f.get('queued') == 1
    ds = summaries[-1] if summaries else {}
    generation = generation if generation is not None else ds.get('generation')
    play, queue = {}, {}
    for line in scoped:
        if '[SPEECH_PLAY]' not in line and '[SPEECH_QUEUE]' not in line:
            continue
        f = fields(line)
        if generation is None or f.get('generation') != generation:
            continue
        if '[SPEECH_PLAY]' in line:
            play = f
        if '[SPEECH_QUEUE]' in line:
            queue = f
    common = sorted(set(chunks) & set(rx))
    common = [s for s in common if chunks[s].get('send_complete_ms') is not None]
    anchor = common[0] if common else None
    rows = []
    sequences = sorted(set(chunks) | set(rx) | (set(range(expected)) if isinstance(expected, int) and 0 <= expected <= 2250 else set()))
    for seq in sequences:
        g, d = chunks.get(seq, {}), rx.get(seq, {})
        gap = None
        if anchor is not None and g.get('send_complete_ms') is not None and d:
            gap = ((d['complete_ms']-rx[anchor]['complete_ms']) -
                   (g['send_complete_ms']-chunks[anchor]['send_complete_ms']))
        rows.append(dict(seq=seq, gateway_send_start=g.get('send_attempt_ms'),
                         gateway_send_complete=g.get('send_complete_ms'), gateway_send_ms=g.get('websocket_send_ms'),
                         gateway_lock_wait_ms=g.get('device_lock_wait_ms'), device_frame_complete=d.get('complete_ms'),
                         network_gap=gap, queue_depth=d.get('queue_depth_after'), played_bytes=d.get('played_bytes'),
                         push=d.get('push')))
    def maximum(key, fallback):
        values = [r[key] for r in sends if isinstance(r.get(key), (float, int))]
        if isinstance(gs.get(fallback), (float, int)):
            values.append(gs[fallback])
        return max(values) if values else None
    max_send = maximum('websocket_send_ms', 'max_websocket_send_ms')
    max_lock = maximum('device_lock_wait_ms', 'max_device_lock_wait_ms')
    max_drift = maximum('drift_from_nominal_ms', 'max_drift_ms')
    first_missing = ds.get('first_missing')
    if first_missing == -1:
        first_missing = None
        # Accepted chunks form a contiguous prefix (firmware SpeechStream validation).
        if isinstance(expected, int) and ds.get('received') == ds.get('accepted') and ds.get('accepted', expected) < expected:
            first_missing = ds['accepted']
    received = ds.get('received')
    end_seen = any(e.get('event') == 'end' for e in events) or ds.get('end_seen') == 1
    receive_timeout = generation is not None and any(
        '[SPEECH] receive idle timeout' in line and fields(line).get('generation') == generation
        for line in scoped)
    # Missing end in a truncated capture is unknown. A matching idle timeout and
    # receive summary prove the device stopped waiting without observing the end.
    end_missing_at_timeout = bool(ds) and receive_timeout and not end_seen
    upstream_complete = bool(gs.get('speech_end_sent')) and expected is not None and gs.get('completed_chunks') == expected
    gateway_normal = (upstream_complete and max_send is not None and max_lock is not None
                      and max_send <= 150 and max_lock <= 100 and (max_drift or 0) <= 150)
    rx_slow = ds.get('max_gap_ms', 0) > 250 or ds.get('max_frame_ms', 0) > 250
    receive_normal = (bool(ds) and not rx_slow and valid_end and received == expected
                      and ds.get('accepted') == expected and first_missing is None)
    classification, reason = 'D', 'Insufficient evidence for A/B/C, or no measured stall.'
    if (max_send or 0) > 150 or (max_lock or 0) > 100:
        classification, reason = 'A', 'Gateway send >150ms or device lock wait >100ms.'
    elif gateway_normal and (rx_slow or first_missing is not None or end_missing_at_timeout):
        classification, reason = 'B', 'Gateway summary normal; device frame/gap >250ms, proven incomplete prefix or end missing at receive timeout.'
    elif gateway_normal and receive_normal and play.get('final') == 1 and (
            play.get('write_errors', 0) > 0 or (expected_bytes is not None and play.get('bytes_written', expected_bytes) < expected_bytes)
            or play.get('max_write_us', 0) > 250000):
        classification, reason = 'C', 'Gateway/receive normal; settled playback has write error, byte deficit or write >250ms.'
    summary = dict(classification=classification, reason=reason, root_cause='not determined',
                   gateway_max_send_ms=max_send, gateway_max_lock_wait_ms=max_lock, gateway_max_drift_ms=max_drift,
                   device_max_receive_gap_ms=ds.get('max_gap_ms'), device_max_frame_ms=ds.get('max_frame_ms'),
                   expected_count=expected, received_count=received, accepted_count=ds.get('accepted'),
                   sampled_sequences=len(rx), first_missing_sequence=first_missing, played_count=queue.get('played'),
                   expected_bytes=expected_bytes, played_bytes=play.get('bytes_written'),
                   playback_final=play.get('final'), speech_end_seen=end_seen if scoped else None,
                   receive_idle_timeout=receive_timeout if scoped else None,
                   speech_end_valid=valid_end if scoped else None,
                   disconnect=any('[CLOUD] disconnected' in l for l in scoped) if scoped else None,
                   malformed_gateway_records=malformed,
                   network_gap_definition='Relative receive-minus-send progression anchored at first common seq; NOT one-way latency.',
                   missing_samples='Blank rows are unobserved samples, not packet loss.',
                   playback_definition='Successful codec API bytes, not proof of DMA drain or audible sound.')
    return dict(speech_id=speech_id, summary=summary, rows=rows)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--gateway', type=Path, required=True)
    parser.add_argument('--serial', type=Path, required=True)
    parser.add_argument('--speech-id', required=True)
    parser.add_argument('--json', type=Path)
    args = parser.parse_args()
    result = correlate(args.gateway.read_text(errors='replace'), args.serial.read_text(errors='replace'), args.speech_id)
    print(json.dumps(result['summary'], indent=2, ensure_ascii=False))
    if result['rows']:
        stream = io.StringIO()
        writer = csv.DictWriter(stream, fieldnames=list(result['rows'][0]))
        writer.writeheader(); writer.writerows(result['rows'])
        print(stream.getvalue(), end='')
    if args.json:
        args.json.write_text(json.dumps(result, indent=2, ensure_ascii=False)+'\n')


if __name__ == '__main__':
    main()
