#!/usr/bin/env python3

import sys
import json
import re
import argparse
from datetime import datetime

def rfc3339_to_epoch(rfc3339_str):
    """Convert RFC 3339 timestamp to Unix epoch seconds.

    Expected format: 2026-08-15T00:00:00.000Z
    Raises ValueError if format is invalid.
    """
    try:
        # Parse RFC3339 UTC timestamp with milliseconds
        dt = datetime.strptime(rfc3339_str, '%Y-%m-%dT%H:%M:%S.%fZ')
        # Convert to epoch seconds (integer)
        epoch_seconds = int(dt.timestamp())
        return epoch_seconds
    except ValueError as e:
        raise ValueError(f"Invalid RFC 3339 timestamp: {str(e)}")

def posix_ere_to_python(posix_ere):
    """Convert POSIX ERE to Python regex.

    Translates POSIX character classes to Python equivalents.
    Raises ValueError if unsupported POSIX constructs are detected.
    """
    posix_translations = {
        '[:alnum:]': 'a-zA-Z0-9',
        '[:alpha:]': 'a-zA-Z',
        '[:digit:]': '0-9',
        '[:lower:]': 'a-z',
        '[:upper:]': 'A-Z',
        '[:space:]': r' \t\n\r\f\v',
        '[:xdigit:]': '0-9a-fA-F',
        '[:punct:]': r'!-/:-@\[-`{-~',
    }

    result = posix_ere

    # Replace POSIX character classes
    for posix_class, python_class in posix_translations.items():
        result = result.replace(posix_class, python_class)

    # Check if there are any remaining [: patterns (unsupported POSIX)
    if '[:' in result and ':]' in result:
        raise ValueError("Unsupported POSIX character class in regex")

    return result

def parse_args():
    parser = argparse.ArgumentParser(
        prog='label-backfill',
        add_help=False,
        usage='label-backfill --room ROOMLOG.jsonl --coder-nick NICK [--model-regex RE] [--legacy-before TS]'
    )

    parser.add_argument('--room', required=True, help='Path to room log file')
    parser.add_argument('--coder-nick', action='append', dest='coder_nicks', required=True,
                       help='Coder nick (repeatable)')
    parser.add_argument('--model-regex', default=r'model ([A-Za-z0-9/_.:-]+)\)',
                       help='Regex to extract model ID')
    parser.add_argument('--legacy-before', default=None,
                       help='RFC 3339 timestamp for legacy cutover')

    args = parser.parse_args()
    return args

def backfill(args):
    room_path = args.room
    coder_nicks = set(args.coder_nicks)
    model_regex = args.model_regex
    legacy_before = args.legacy_before

    # Convert POSIX ERE to Python regex
    try:
        python_regex = posix_ere_to_python(model_regex)
    except ValueError as e:
        print(f"label-backfill: error: {str(e)}", file=sys.stderr)
        sys.exit(2)

    # Compile regex
    try:
        compiled_regex = re.compile(python_regex)
    except re.error as e:
        print(f"label-backfill: error: invalid regex: {str(e)}", file=sys.stderr)
        sys.exit(2)

    # Convert legacy_before from RFC3339 to epoch seconds if provided
    legacy_before_epoch = None
    if legacy_before:
        try:
            legacy_before_epoch = rfc3339_to_epoch(legacy_before)
        except ValueError as e:
            print(f"label-backfill: error: {str(e)}", file=sys.stderr)
            sys.exit(2)

    # Read room log
    try:
        with open(room_path, 'r', encoding='utf-8') as f:
            for lineno, line in enumerate(f, 1):
                line = line.rstrip('\n\r')

                # Parse JSON
                try:
                    msg = json.loads(line)
                except (json.JSONDecodeError, ValueError):
                    print(f"label-backfill: error: line {lineno}: invalid JSON", file=sys.stderr)
                    sys.exit(1)

                # Validate required fields (verified, in_reply_to, and unknown keys are optional/tolerated)
                required_fields = {'id', 'ts', 'from', 'to', 'type', 'body'}
                if not isinstance(msg, dict) or not required_fields.issubset(msg.keys()):
                    print(f"label-backfill: error: line {lineno}: missing required fields", file=sys.stderr)
                    sys.exit(1)

                msg_id = msg['id']
                ts = msg['ts']
                from_nick = msg['from']
                body = msg['body']

                # Determine model_id and reason
                model_id = None
                reason = None

                if from_nick not in coder_nicks:
                    # Not a coder - no model
                    reason = "role_no_model"
                else:
                    # Is a coder - try to extract model from body
                    match = compiled_regex.search(body)
                    if match:
                        # Regex matched - extract model_id
                        model_id = match.group(1)
                        reason = None
                    else:
                        # Regex didn't match
                        if legacy_before_epoch and ts < legacy_before_epoch:
                            # Before legacy cutover
                            reason = "legacy_pre_schema"
                        else:
                            # Parse failed
                            reason = "parse_failed"

                # Output mapping
                output = {
                    "msg_id": msg_id,
                    "model_id": model_id,
                    "model_id_reason": reason
                }
                print(json.dumps(output, separators=(',', ':'), ensure_ascii=True))

    except (IOError, OSError) as e:
        print(f"label-backfill: error: cannot read file: {str(e)}", file=sys.stderr)
        sys.exit(2)

    return 0

if __name__ == '__main__':
    try:
        args = parse_args()
    except SystemExit as e:
        if e.code != 0:
            print("label-backfill: error: missing or invalid arguments", file=sys.stderr)
            sys.exit(2)
        sys.exit(e.code)

    sys.exit(backfill(args))
