#!/usr/bin/env python3
"""
Audit tool for C-1c contract: verify correlation_id consistency.

Usage: corr-audit --room ROOMLOG.jsonl --manifest MANIFEST.jsonl [--since TS] [--require-spawned]
  Checks asks (type=="ask", ts >= since) and manifest lines for correlation consistency.
"""

import sys
import json
import re
import argparse
from datetime import datetime, timezone


# UUID v4 ERE from C-1c §2.1
UUID_RE = r'^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$'


def is_valid_correlation_id(value):
    """Check if value matches the UUID v4 ERE."""
    if not isinstance(value, str):
        return False
    return bool(re.match(UUID_RE, value))


def parse_rfc3339_ts(ts_str):
    """Parse RFC 3339 timestamp and return as Unix epoch seconds (integer)."""
    try:
        # Handle both microsecond and millisecond precision
        if '.' in ts_str:
            # Has fractional seconds
            if ts_str.endswith('Z'):
                ts_str = ts_str[:-1] + '+00:00'
            dt = datetime.fromisoformat(ts_str)
        else:
            if ts_str.endswith('Z'):
                ts_str = ts_str[:-1] + '+00:00'
            dt = datetime.fromisoformat(ts_str)
        # Convert to Unix epoch seconds (integer)
        return int(dt.timestamp())
    except ValueError:
        return None


def main():
    parser = argparse.ArgumentParser(
        prog='corr-audit',
        description='Audit correlation_id consistency across room and manifest'
    )
    parser.add_argument('--room', dest='room_file', required=True,
                        help='Room log JSONL file')
    parser.add_argument('--manifest', dest='manifest_file', required=True,
                        help='Manifest JSONL file')
    parser.add_argument('--since', dest='since_ts', default=None,
                        help='RFC 3339 UTC timestamp: exempt asks with ts < since')
    parser.add_argument('--require-spawned', action='store_true',
                        help='Fail on unspawned asks (default: warn)')

    try:
        args = parser.parse_args()
    except SystemExit as e:
        if e.code != 0:
            sys.exit(2)
        raise

    # Parse --since if provided (convert RFC3339 to Unix epoch seconds)
    since_epoch = None
    if args.since_ts:
        since_epoch = parse_rfc3339_ts(args.since_ts)
        if since_epoch is None:
            print(f'corr-audit: error: invalid --since timestamp: {args.since_ts}',
                  file=sys.stderr)
            sys.exit(2)

    # Read room log
    asks = []
    try:
        with open(args.room_file, 'r') as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    entry = json.loads(line)
                    # Skip non-object lines (arrays, primitives, etc.)
                    if not isinstance(entry, dict):
                        continue
                    if entry.get('type') == 'ask':
                        # Check if this ask is in scope (ts is now a Unix epoch integer)
                        ts = entry.get('ts')
                        if ts and since_epoch is not None:
                            # ts should be an integer (Unix epoch seconds)
                            if isinstance(ts, int) and ts < since_epoch:
                                # Pre-cutover, skip
                                continue
                        asks.append(entry)
                except json.JSONDecodeError:
                    pass
    except IOError as e:
        print(f'corr-audit: error: cannot read room log: {e}', file=sys.stderr)
        sys.exit(2)

    # Read manifest
    manifest_lines = []
    try:
        with open(args.manifest_file, 'r') as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    entry = json.loads(line)
                    # Skip non-object lines (arrays, primitives, etc.)
                    if not isinstance(entry, dict):
                        continue
                    manifest_lines.append(entry)
                except json.JSONDecodeError:
                    pass
    except IOError as e:
        print(f'corr-audit: error: cannot read manifest: {e}', file=sys.stderr)
        sys.exit(2)

    # Run checks
    findings = []  # (severity, code, entity_type, entity_id)
    warnings = []

    # Track state
    ask_ids_with_cid = {}  # correlation_id -> [ask_id, ...]
    manifest_cids = {}  # correlation_id -> [session_id, ...]
    failed_asks = set()  # asks that failed A or B

    # Check A: MISSING correlation_id
    for ask in asks:
        ask_id = ask.get('id')
        if 'correlation_id' not in ask:
            findings.append(('FAIL', 'MISSING', 'ask', ask_id))
            failed_asks.add(ask_id)

    # Check B: MALFORMED correlation_id
    for ask in asks:
        ask_id = ask.get('id')
        if ask_id in failed_asks:
            continue  # Already failed A
        if 'correlation_id' in ask:
            cid = ask['correlation_id']
            if not is_valid_correlation_id(cid):
                findings.append(('FAIL', 'MALFORMED', 'ask', ask_id))
                failed_asks.add(ask_id)

    # Build maps for C, D, E checks
    for ask in asks:
        ask_id = ask.get('id')
        if ask_id in failed_asks:
            continue
        if 'correlation_id' in ask:
            cid = ask['correlation_id']
            if is_valid_correlation_id(cid):
                if cid not in ask_ids_with_cid:
                    ask_ids_with_cid[cid] = []
                ask_ids_with_cid[cid].append(ask_id)

    for entry in manifest_lines:
        cid = entry.get('correlation_id')
        if cid is not None:  # Only non-null
            if cid not in manifest_cids:
                manifest_cids[cid] = []
            manifest_cids[cid].append(entry.get('session_id'))

    # Check C: DUP_ASK
    for cid, ask_ids in ask_ids_with_cid.items():
        if len(ask_ids) > 1:
            # Report each extra ask
            for ask_id in ask_ids[1:]:
                findings.append(('FAIL', 'DUP_ASK', 'ask', ask_id))

    # Check D: DUP_SESSION
    for cid, session_ids in manifest_cids.items():
        if len(session_ids) > 1:
            # Report each extra session
            for session_id in session_ids[1:]:
                findings.append(('FAIL', 'DUP_SESSION', 'session', session_id))

    # Check E: ORPHAN_SESSION
    ask_cids = set(ask_ids_with_cid.keys())
    for cid, session_ids in manifest_cids.items():
        if cid not in ask_cids:
            # This manifest cid has no matching ask
            for session_id in session_ids:
                findings.append(('FAIL', 'ORPHAN_SESSION', 'session', session_id))

    # Check F: UNSPAWNED
    # An ask is unspawned if its correlation_id doesn't appear in manifest
    # Excluded: asks that failed A or B
    manifest_all_cids = set()
    for entry in manifest_lines:
        if entry.get('correlation_id') is not None:
            manifest_all_cids.add(entry.get('correlation_id'))

    for ask in asks:
        ask_id = ask.get('id')
        if ask_id in failed_asks:
            continue  # A or B failed, exclude from F
        if 'correlation_id' in ask:
            cid = ask['correlation_id']
            if is_valid_correlation_id(cid):
                if cid not in manifest_all_cids:
                    if args.require_spawned:
                        findings.append(('FAIL', 'UNSPAWNED', 'ask', ask_id))
                    else:
                        warnings.append(('WARN', 'UNSPAWNED', 'ask', ask_id))

    # Output findings
    fail_count = 0
    warn_count = 0

    for severity, code, entity_type, entity_id in findings:
        if severity == 'FAIL':
            if entity_type == 'ask':
                print(f'FAIL {code} ask={entity_id}')
            else:
                print(f'FAIL {code} session={entity_id}')
            fail_count += 1

    for severity, code, entity_type, entity_id in warnings:
        if severity == 'WARN':
            if entity_type == 'ask':
                print(f'WARN {code} ask={entity_id}')
            else:
                print(f'WARN {code} session={entity_id}')
            warn_count += 1

    # Summary line
    num_asks = len(asks)
    num_sessions = len(manifest_lines)
    print(f'corr-audit: asks={num_asks} sessions={num_sessions} fail={fail_count} warn={warn_count}')

    # Exit code
    if fail_count > 0:
        sys.exit(1)
    else:
        sys.exit(0)


if __name__ == '__main__':
    main()
