#!/usr/bin/env python
import argparse
import json
import re
from collections import defaultdict
from pathlib import Path

import pyarrow.parquet as pq

from eval_segmentqa_vllm import correct


BUCKETS = {'1': 'or', '2': 'tg', '3': 'agg', '4': 'eu', '5': 'cr'}
TIME_RE = re.compile(r'\d{2}:\d{2}:\d{2}(?:\s*,\s*\d{2}:\d{2}:\d{2})*')


def seconds(value):
    hours, minutes, secs = map(int, value.split(':'))
    if minutes >= 60 or secs >= 60:
        raise ValueError(value)
    return hours * 3600 + minutes * 60 + secs


def parse_times(value):
    value = re.sub(r'<think>.*?</think>', '', str(value or ''), flags=re.S).strip()
    if not TIME_RE.fullmatch(value):
        return None
    try:
        return sorted(seconds(item.strip()) for item in value.split(','))
    except ValueError:
        return None


def summarize(decisions, total):
    by_bucket = {}
    for bucket in ('agg', 'tg', 'or', 'eu', 'cr'):
        values = decisions[bucket]
        passed = sum(values)
        by_bucket[bucket] = {'correct': passed, 'total': len(values),
                             'accuracy': passed / len(values) if values else 0.0}
    passed = sum(sum(values) for values in decisions.values())
    return {'correct': passed, 'total': total, 'accuracy': passed / total if total else 0.0,
            'by_bucket': by_bucket}


def summarize_formats(decisions):
    by_format = {}
    for fmt, values in decisions.items():
        passed = sum(values)
        by_format[fmt] = {'correct': passed, 'total': len(values),
                          'accuracy': passed / len(values) if values else 0.0}
    return by_format


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--predictions', required=True)
    parser.add_argument('--references', required=True)
    parser.add_argument('--parquet', action='append', required=True)
    parser.add_argument('--local-output', required=True)
    parser.add_argument('--time-output', required=True)
    args = parser.parse_args()

    predictions = [json.loads(line) for line in Path(args.predictions).read_text().splitlines() if line.strip()]
    references = {str(row['id']): row for row in
                  (json.loads(line) for line in Path(args.references).read_text().splitlines() if line.strip())}
    metadata = {}
    for path in args.parquet:
        metadata.update({str(row['id']): row for row in pq.read_table(path).to_pylist()})

    ids = [str(row['id']) for row in predictions]
    expected = set(references)
    actual = set(ids)
    completeness = {
        'predictions': len(predictions),
        'unique_ids': len(actual),
        'missing': len(expected - actual),
        'duplicate': len(ids) - len(actual),
        'unexpected': len(actual - expected),
        'failed': sum(bool(row.get('error') or row.get('failed')) for row in predictions),
        'skipped': sum(bool(row.get('skipped')) for row in predictions),
        'empty_predictions': sum(not str(row.get('prediction', '')).strip() for row in predictions),
    }
    if completeness != {'predictions': len(expected), 'unique_ids': len(expected), 'missing': 0,
                         'duplicate': 0, 'unexpected': 0, 'failed': 0, 'skipped': 0,
                         'empty_predictions': 0}:
        raise RuntimeError(f'incomplete predictions: {completeness}')
    if set(metadata) != expected:
        raise RuntimeError(f'official parquet/reference ID mismatch: parquet={len(metadata)} refs={len(expected)}')

    local, tolerant = defaultdict(list), defaultdict(list)
    local_formats, tolerant_formats = defaultdict(list), defaultdict(list)
    time_parse_errors = 0
    for row in predictions:
        qid = str(row['id'])
        reference, meta = references[qid], metadata[qid]
        bucket = BUCKETS[str(meta['primary_capability'])[0]]
        local_result = correct(row['prediction'], reference['answer'], reference['answer_format'])
        local[bucket].append(local_result)
        local_formats[reference['answer_format']].append(local_result)
        tolerant_result = local_result
        if reference['answer_format'] == 'time':
            pred_times, gold_times = parse_times(row['prediction']), parse_times(reference['answer'])
            if pred_times is None:
                time_parse_errors += 1
                tolerant_result = False
            else:
                duration = seconds(meta['timestamp_end']) - seconds(meta['timestamp_start'])
                tolerance = min(5.0, 1.0 + duration * 4.0 / 360.0)
                tolerant_result = len(pred_times) == len(gold_times) and all(
                    abs(pred - gold) <= tolerance for pred, gold in zip(pred_times, gold_times))
        tolerant[bucket].append(tolerant_result)
        tolerant_formats[reference['answer_format']].append(tolerant_result)

    local_metrics = {'protocol': 'scripts/eval_segmentqa_vllm.py::correct',
                     **completeness, **summarize(local, len(predictions)),
                     'by_format': summarize_formats(local_formats)}
    tolerant_metrics = {
        'protocol': 'time-tolerance-only',
        'time_tolerance_seconds': 'min(5, 1 + segment_duration * 4 / 360)',
        'non_time_protocol': 'scripts/eval_segmentqa_vllm.py::correct',
        'time_parse_errors': time_parse_errors,
        **completeness,
        **summarize(tolerant, len(predictions)),
        'by_format': summarize_formats(tolerant_formats),
    }
    Path(args.local_output).write_text(json.dumps(local_metrics, indent=2))
    Path(args.time_output).write_text(json.dumps(tolerant_metrics, indent=2))
    print(json.dumps({'local': local_metrics, 'time_tolerance': tolerant_metrics}, indent=2))


if __name__ == '__main__':
    main()