#!/usr/bin/env python3
"""Merge permanent rejects from a local live log into baked marks."""
import argparse
import datetime
import json
from pathlib import Path
import re
import sys

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from src.config_params import GPU_ENUMS
from src.live_queue import model_key
from src.marks import PERMANENT_PAIR_CODES, REJECT_CODES, load_marks


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('log')
    parser.add_argument('--marks', default='src/data/reject-marks.json')
    parser.add_argument('--date', default=datetime.date.today().isoformat())
    parser.add_argument('--dry-run', action='store_true')
    args = parser.parse_args()
    path = Path(args.marks)
    if path.exists():
        load_marks(path)
        rows = json.loads(path.read_text())['marks']
    else:
        rows = []
    merged = {}
    for row in rows:
        merged.setdefault((model_key(row['model']), row['gpu']), row)
    before = len(merged)
    for line in Path(args.log).read_text().splitlines():
        match = re.search(r'\[live\] (rejected|skip) (.*)', line)
        if not match:
            continue
        fields = dict(token.split('=', 1) for token in match[2].split() if '=' in token)
        model, gpu = fields.get('model', ''), fields.get('gpu', '')
        raw = fields.get('code', '') if match[1] == 'rejected' else fields.get('reason', '').removeprefix('marked:precheck_')
        if match[1] == 'skip' and not fields.get('reason', '').startswith('marked:precheck_'):
            continue
        try:
            code = int(raw)
        except ValueError:
            continue
        if (code not in PERMANENT_PAIR_CODES or gpu not in GPU_ENUMS.values()
                or not re.fullmatch(r'[\w.-]+/[\w.\-]+', model)):
            continue
        merged.setdefault((model_key(model), gpu), dict(model=model, gpu=gpu, code=code,
                          reason=REJECT_CODES[code], date=args.date))
    rows = sorted(merged.values(), key=lambda r: (r['model'].lower(), r['gpu']))
    if not args.dry_run:
        path.write_text(json.dumps(dict(version=1, marks=rows), indent=2) + '\n')
    print(f'added {len(rows) - before} total {len(rows)}')


if __name__ == '__main__':
    main()