import argparse
import json
import sys
from collections import Counter
from pathlib import Path


def spans(tokens, labels):
    i = 0
    n = min(len(tokens), len(labels))
    while i < n:
        lab = labels[i]
        if not isinstance(lab, str) or not lab.startswith('B-'):
            i += 1
            continue
        ent = lab[2:]
        start = i
        i += 1
        while i < n and labels[i] == f'I-{ent}':
            i += 1
        end = i
        text = ''.join(tokens[start:end])
        yield ent, start, end, text


def normalize_season_text(text: str) -> str:
    t = (text or '').strip()
    if not t:
        return t
    return ' '.join(t.split())


def main() -> None:
    parser = argparse.ArgumentParser(description='Analyze season label distribution from JSONL char labels.')
    parser.add_argument('--input', required=True)
    parser.add_argument('--top-k', type=int, default=30)
    parser.add_argument('--sample-limit', type=int, default=8)
    parser.add_argument('--output', default=None)
    args = parser.parse_args()

    path = Path(args.input)
    season_counter = Counter()
    season_examples = {}
    weird_patterns = {
        'roman': 0,
        'ni_like': 0,
        'cjk_ord': 0,
        's_prefix': 0,
        'english_ord': 0,
        'digits_only': 0,
        'other': 0,
    }

    target_rows = {
        'kakuriyo_ni': [],
        'fire_force_nino': [],
    }

    rows = 0
    season_rows = 0

    with path.open('r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            rows += 1
            obj = json.loads(line)
            tokens = obj.get('tokens') or []
            labels = obj.get('labels') or []
            filename = obj.get('filename') or ''
            lower = filename.lower()

            if 'kakuriyo no yadomeshi ni' in lower and len(target_rows['kakuriyo_ni']) < args.sample_limit:
                target_rows['kakuriyo_ni'].append({
                    'filename': filename,
                    'season_field': obj.get('season'),
                    'title_field': obj.get('title'),
                })
            if (
                '炎炎' in filename
                and any(key in filename for key in ['弐ノ章', '貳ノ章', '貳之章'])
                and len(target_rows['fire_force_nino']) < args.sample_limit
            ):
                target_rows['fire_force_nino'].append({
                    'filename': filename,
                    'season_field': obj.get('season'),
                    'title_field': obj.get('title'),
                })

            found = False
            for ent, _s, _e, text in spans(tokens, labels):
                if ent != 'SEASON':
                    continue
                found = True
                season_rows += 1
                season = normalize_season_text(text)
                season_counter[season] += 1
                if season not in season_examples:
                    season_examples[season] = filename

                upper = season.upper()
                if upper in {'I', 'II', 'III', 'IV', 'V', 'VI'}:
                    weird_patterns['roman'] += 1
                elif season.lower().startswith('ni'):
                    weird_patterns['ni_like'] += 1
                elif any(ch in season for ch in ['第', '季', '章', '部', '期', '弐', '貳', '二']):
                    weird_patterns['cjk_ord'] += 1
                elif upper.startswith('S') and len(upper) >= 2 and upper[1].isdigit():
                    weird_patterns['s_prefix'] += 1
                elif any(x in season.lower() for x in ['season', 'st', 'nd', 'rd', 'th']):
                    weird_patterns['english_ord'] += 1
                elif season.isdigit():
                    weird_patterns['digits_only'] += 1
                else:
                    weird_patterns['other'] += 1

            if not found:
                pass

    result = {
        'input': str(path),
        'rows': rows,
        'season_span_count': int(sum(season_counter.values())),
        'unique_season_values': len(season_counter),
        'top_season_values': [
            {
                'season': season,
                'count': count,
                'percent_of_season_spans': (count / max(1, sum(season_counter.values()))) * 100.0,
                'example_filename': season_examples.get(season),
            }
            for season, count in season_counter.most_common(args.top_k)
        ],
        'focus_values': {
            key: season_counter.get(key, 0)
            for key in [
                'Ni',
                'ni',
                'NI',
                'II',
                '2',
                '02',
                'S2',
                'S02',
                'Season 2',
                '2nd',
                '第2季',
                '第二季',
                '弐ノ章',
                '貳ノ章',
                '貳之章',
                '貳',
            ]
        },
        'pattern_buckets': weird_patterns,
        'target_rows': target_rows,
    }

    text = json.dumps(result, ensure_ascii=False, indent=2)
    if args.output:
        out = Path(args.output)
        out.parent.mkdir(parents=True, exist_ok=True)
        out.write_text(text + '\n', encoding='utf-8')
        print(f'WROTE {out}')
    try:
        print(text)
    except UnicodeEncodeError:
        # Some Windows consoles still default to GBK. Emit UTF-8 bytes directly.
        sys.stdout.buffer.write((text + '\n').encode('utf-8', errors='replace'))


if __name__ == '__main__':
    main()