"""Document rewriting with the desktop generation policy and incremental reports."""
import copy
import sys
import threading
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent / 'engine'))
import documents, generation, geometry, rewrite_policy


def transform(engine, text, g, style, options, output, cancel=None):
    cancel = cancel or threading.Event()
    locations = documents.spans(text, limit=700)
    records = {}
    report = dict(policy=rewrite_policy.VERSION, identity=engine.identity,
                  input=text, geometry=g, style=style, sampling=options,
                  spans=locations, records=records, complete=False)
    output = Path(output)

    def save():
        report['text'] = documents.assemble(text, locations, records)
        geometry.atomic(output, report)
        output.with_suffix('.txt').write_text(report['text'], encoding='utf-8')

    if output.suffix.lower() != '.json':
        raise ValueError('The report output must end in .json; rewritten text uses the matching .txt path.')
    save()
    for index, (a, b) in enumerate(locations):
        if cancel.is_set():
            break
        engine.progress(f'Rewriting passage {index + 1}/{len(locations)}')
        limit = documents.rewrite_allowance(text, locations, records, index, multiplier=1.25)
        target = geometry.point(g, (a + b) / (2 * len(text)))
        instruction = rewrite_policy.prompt(style) + (
            f'\nAim for at most {b-a} characters; finish before the hard limit of '
            f'{limit} characters, including spaces and punctuation.')
        opts = copy.deepcopy(options)
        opts['max_tokens'] = min(2048, max(opts['max_tokens'], limit))

        def generate(adjusted, attempt):
            return engine.generate(
                text[a:b], g=adjusted, coordinates=target, system=instruction,
                max_chars=limit, soft_chars=max(1, int((b-a)*.85)),
                cancel=cancel, **opts)
        try:
            result = generation.rewrite(text[a:b], generate, g, True, cancel,
                                        engine.progress, max_chars=limit)
        except InterruptedError:
            break
        records[str(index)] = dict(result, target=target, character_allowance=limit)
        save()
    report['complete'] = len(records) == len(locations)
    report['failed'] = sum(bool(r.get('original_retained')) for r in records.values())
    report['needs_review'] = sum(bool(r.get('warning') or r.get('error')) for r in records.values())
    save()
    return report