"""Classify long text in token windows, preserving per-window evidence."""
import argparse
import json
import os
import re
from pathlib import Path
os.environ.setdefault('HF_HOME', str(Path(__file__).resolve().parent / '.cache/huggingface'))

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

ROOT = Path(__file__).resolve().parent
TASKS = ['open-emotions', 'open-morals', 'open-sensations']


def make_windows(tokenizer, text, window_size, overlap, target=None):
    if not text.strip():
        raise ValueError('Text must not be empty.')
    prefix = f'Target: {target}\nContext: ' if target is not None else ''
    # Tokenize the full document, but only feed bounded slices to the model.
    encoded = tokenizer(prefix + text, add_special_tokens=False, return_offsets_mapping=True, verbose=False)
    offsets = encoded['offset_mapping']
    backend_encoding = None
    backend_state = None
    if not hasattr(tokenizer, 'prepare_for_model'):
        # Transformers 5 delegates special-token processing to the Rust backend.
        backend_encoding = tokenizer.backend_tokenizer.encode(prefix + text, add_special_tokens=False)
        if backend_encoding.ids != encoded['input_ids']:
            raise ValueError('Backend tokenization differs from window offsets.')
        backend_state = json.loads(backend_encoding.__getstate__())
    prefix_count = sum(end <= len(prefix) for _, end in offsets) if prefix else 0
    prefix_ids = encoded['input_ids'][:prefix_count]
    ids = encoded['input_ids'][prefix_count:]
    offsets = offsets[prefix_count:]
    capacity = window_size - tokenizer.num_special_tokens_to_add(pair=False) - len(prefix_ids)
    if capacity <= 0 or not 0 <= overlap < capacity:
        raise ValueError(f'Overlap must be between 0 and {capacity - 1}; prefix and special tokens consume the rest of the window.')
    mentions = list(re.finditer(r'(?<!\w)' + re.escape(target) + r'(?!\w)', text, re.IGNORECASE)) if target else []
    if target is not None and not mentions:
        raise ValueError('The sensory target must occur in the context text.')
    windows = []
    skipped = 0
    start = 0
    while start < len(ids):
        end = min(start + capacity, len(ids))
        char_start = max(0, offsets[start][0] - len(prefix))
        char_end = offsets[end - 1][1] - len(prefix)
        contains_target = target is None or any(char_start <= m.start() and m.end() <= char_end for m in mentions)
        if contains_target:
            if backend_encoding is None:
                prepared = tokenizer.prepare_for_model(prefix_ids + ids[start:end], add_special_tokens=True,
                                                       return_attention_mask=True, truncation=False)
            else:
                def slice_encoding(left, right):
                    # Reconstruct through Encoding's pickle state. Native truncate
                    # retains overflow chunks; merging those creates cross-products.
                    state = dict(backend_state)
                    for field in ['ids', 'type_ids', 'tokens', 'words', 'offsets',
                                  'special_tokens_mask', 'attention_mask']:
                        state[field] = backend_state[field][left:right]
                    state['overflowing'] = []
                    state['sequence_ranges'] = {'0': {'start': 0, 'end': right - left}}
                    segment = type(backend_encoding)()
                    segment.__setstate__(json.dumps(state).encode())
                    return segment
                segment = slice_encoding(prefix_count + start, prefix_count + end)
                if prefix_count:
                    segment = type(segment).merge([slice_encoding(0, prefix_count), segment], growing_offsets=False)
                prepared_encoding = tokenizer.backend_tokenizer.post_process(segment, add_special_tokens=True)
                prepared = dict(input_ids=prepared_encoding.ids, attention_mask=prepared_encoding.attention_mask)
            if len(prepared['input_ids']) > window_size:
                raise ValueError('Window exceeds the requested limit.')
            windows.append(dict(token_start=start, token_end=end, char_start=char_start, char_end=char_end,
                                input_tokens=len(prepared['input_ids']), text=text[char_start:char_end], encoded=prepared))
        else:
            skipped += 1
        if end == len(ids):
            break
        start += capacity - overlap
    if not windows:
        raise ValueError('No window contains the complete target. Increase window size or overlap.')
    return windows, dict(context_tokens=len(ids), context_capacity=capacity, overlap_tokens=overlap,
                         skipped_windows_without_complete_target=skipped)


def classify_windows(model, tokenizer, windows, threshold, batch_size=8):
    rows = []
    for start in range(0, len(windows), batch_size):
        batch = tokenizer.pad([w['encoded'] for w in windows[start:start + batch_size]],
                              padding=True, return_tensors='pt')
        with torch.inference_mode():
            rows.extend(model(**batch).logits.sigmoid().cpu().tolist())
    labels = [model.config.id2label[i] for i in range(len(rows[0]))]
    result = []
    for window, row in zip(windows, rows):
        item = {k: v for k, v in window.items() if k != 'encoded'}
        item.update(labels=[label for label, score in zip(labels, row) if score >= threshold],
                    scores={label: round(score, 6) for label, score in zip(labels, row)})
        result.append(item)
    # These summaries have no separately calibrated document-level decision threshold.
    summaries = {label: dict(max_score=round(max(row[i] for row in rows), 6),
                            mean_score=round(sum(row[i] for row in rows) / len(rows), 6))
                 for i, label in enumerate(labels)}
    return result, summaries


def load_task(task, override=None):
    if override:
        path = override
    elif task == 'open-emotions':
        path = ROOT / 'outputs/open-emotions/model'
    else:
        out = ROOT / 'outputs' / task
        decision = json.loads((out / 'decision.json').read_text(encoding='utf-8'))
        path = out / decision['selected_arm'] / 'model'
    options = json.loads((path / 'thresholds.json').read_text(encoding='utf-8'))
    run = json.loads((path / 'training_run.json').read_text(encoding='utf-8'))
    if not run.get('completed') or run.get('mode') != 'full':
        raise ValueError('A completed full-training checkpoint is required.')
    if run.get('task', 'open-emotions') != task:
        raise ValueError('Checkpoint task does not match the requested classifier.')
    if (path / 'classifier_config.json').exists():
        from gemma_classifier import load_classifier
        model, tokenizer, options = load_classifier(path, device='cuda' if torch.cuda.is_available() else 'cpu')
        return model.eval(), tokenizer, options
    tokenizer = AutoTokenizer.from_pretrained(path)
    model = AutoModelForSequenceClassification.from_pretrained(path, attn_implementation='eager').eval()
    return model, tokenizer, options


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--task', required=True, choices=TASKS)
    inputs = parser.add_mutually_exclusive_group(required=True)
    inputs.add_argument('--text')
    inputs.add_argument('--file', type=Path, help='UTF-8 text file')
    parser.add_argument('--target', help='Required word appearing in sensory context')
    parser.add_argument('--model', type=Path)
    parser.add_argument('--mode', choices=['chunks', 'sliding'], default='sliding')
    parser.add_argument('--window-size', type=int, help='Total tokens including special tokens and sensory prefix')
    parser.add_argument('--overlap', type=int, help='Context tokens shared by adjacent windows')
    parser.add_argument('--output', type=Path)
    args = parser.parse_args()
    if args.task == 'open-sensations' and not args.target:
        parser.error('Open Sensations requires --target WORD')
    if args.task != 'open-sensations' and args.target:
        parser.error('--target applies only to Open Sensations')
    torch.set_num_threads(8)
    model, tokenizer, options = load_task(args.task, args.model)
    size = options['max_length'] if args.window_size is None else args.window_size
    if not 4 <= size <= options['max_length']:
        parser.error(f'Window size must be 4..{options["max_length"]}; stay within the evaluated input limit.')
    overlap = 0 if args.mode == 'chunks' else (args.overlap if args.overlap is not None else size // 4)
    if args.mode == 'chunks' and args.overlap not in (None, 0):
        parser.error('Chunks use zero overlap; select sliding for overlapping windows.')
    text = args.file.read_text(encoding='utf-8') if args.file else args.text
    try:
        windows, metadata = make_windows(tokenizer, text, size, overlap, args.target)
    except ValueError as error:
        parser.error(str(error))
    evidence, summaries = classify_windows(model, tokenizer, windows, options['global_threshold'])
    result = dict(task=args.task, target=args.target, mode=args.mode, window_size=size,
                  trained_input_limit=options['max_length'], threshold=options['global_threshold'],
                  **metadata, windows=evidence, exploratory_summary=summaries,
                  interpretation='Per-window labels use the existing short-input threshold. Max/mean summaries are uncalibrated; overlap repeats evidence. Sensory windows score target associations, not literal sensation presence.')
    payload = json.dumps(result, indent=2, ensure_ascii=False) + '\n'
    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(payload, encoding='utf-8')
    else:
        print(payload)


if __name__ == '__main__':
    main()