"""A bounded tool message with one copy of evidence and no retrieval diagnostics."""
from __future__ import annotations

import copy
import json
import os

from tools.rag.build_chunks import TokenCounter
from tools.text_budget import split_to_fit

SOURCE_FIELDS = ('source_id', 'title', 'author', 'document_path', 'source_url', 'published',
                 'categories', 'tags', 'section')
FIELDS = ('status', 'reason', 'query', 'engine', 'strategies', 'instructions', 'answer_format',
          'complete', 'truncated', 'truncation_reasons', 'index_version', 'source_count')


def compact_rag_payload(data: dict, *, max_tokens: int | None = None,
                        include_diagnostics: bool = False) -> dict:
    from tools.rag.rag_utils import build_context_text

    limit = max_tokens if max_tokens is not None else int(os.getenv('RAG_LLM_MAX_TOKENS', '12000'))
    if limit < 256:
        raise ValueError('RAG_LLM_MAX_TOKENS musí byť aspoň 256')
    encoding = os.getenv('RAG_LLM_TOKEN_ENCODING', 'cl100k_base')
    counter = TokenCounter(encoding)
    if counter.encoding is None:
        raise RuntimeError('Kontextový rozpočet vyžaduje platný tiktoken tokenizér')
    base = {key: copy.deepcopy(data[key]) for key in FIELDS if key in data}
    base.setdefault('status', 'ok')
    base.setdefault('complete', None)
    base.setdefault('truncated', False)
    base.setdefault('truncation_reasons', [])
    input_sources = [{**source, 'source_id': source.get('source_id') or f'S{number}'}
                     for number, source in enumerate(data.get('sources', []), 1)]
    references = [{key: source[key] for key in SOURCE_FIELDS if key in source}
                  for source in input_sources]

    def stamp(payload):
        payload['context_budget'] = {'max_tokens': limit, 'used_tokens': 0,
            'tokenizer': counter.name, 'tokens_are_estimate': True}
        for _ in range(8):
            used = counter.count(json.dumps(payload, ensure_ascii=False, separators=(',', ':')))
            if payload['context_budget']['used_tokens'] == used:
                break
            payload['context_budget']['used_tokens'] = used
        return payload

    def render(sources, *, truncated=False, raw_context=None):
        payload = copy.deepcopy(base)
        if truncated:
            payload['complete'] = False
            payload['truncated'] = True
            payload['truncation_reasons'] = list(dict.fromkeys(
                [*payload['truncation_reasons'], 'llm_context_budget']))
        payload['sources'] = [{key: source[key] for key in SOURCE_FIELDS if key in source} for source in sources]
        payload['source_count'] = len(sources)
        coverage = {key: payload.get(key) for key in ['complete', 'truncated', 'truncation_reasons', 'index_version']}
        context = raw_context if raw_context is not None else build_context_text(sources)
        payload['context'] = 'METADÁTA VYHĽADANIA\n' + json.dumps(coverage, ensure_ascii=False) + '\n\n' + context
        return stamp(payload)

    def fits(payload):
        return payload['context_budget']['used_tokens'] <= limit

    # An already projected server response needs no reconstruction by its client.
    if data.get('context_budget'):
        result = stamp({**base, 'sources': references, 'context': data.get('context', '')})
        if not fits(result):
            parts = split_to_fit(data.get('context', ''), lambda text:
                fits(render(references, truncated=True, raw_context=text)))
            result = render(references, truncated=True, raw_context=parts[0])
    elif any('text' in source for source in input_sources):
        sources = input_sources
        result = render(sources)
        if not fits(result):
            selected = []
            if not fits(render([], truncated=True)):
                raise ValueError('Tokenový rozpočet je menší ako povinné pokyny a metadáta')
            for source in sources:
                candidate = render([*selected, source], truncated=True)
                if fits(candidate):
                    selected.append(source)
                    continue
                body = source.get('text') or ''
                suffix = '\n[OBSAH ZDROJA OREZANÝ]'
                def body_fits(text):
                    return fits(render([*selected, {**source, 'text': text + suffix}], truncated=True))
                if body and body_fits(body[:1]):
                    # Only the first fitting prefix is delivered; the complete
                    # original remains available through the diagnostics option.
                    low, high = 1, len(body)
                    while low < high:
                        middle = (low + high + 1) // 2
                        if body_fits(body[:middle]): low = middle
                        else: high = middle - 1
                    selected.append({**source, 'text': body[:low] + suffix})
                break
            result = render(selected, truncated=True)
            if not selected:
                result['status'] = 'no_answer'
                result['reason'] = 'llm_context_budget_exhausted'
                result = stamp(result)
    else:
        # Compatibility with older tool servers that expose only context + refs.
        context = data.get('context', '')
        result = render(references, raw_context=context)
        if not fits(result):
            parts = split_to_fit(context, lambda text: fits(render(references, truncated=True, raw_context=text)))
            result = render(references, truncated=True, raw_context=parts[0])
    if not fits(result):
        raise ValueError('Povinné metadáta prekračujú kontextový rozpočet')
    if include_diagnostics:
        result['diagnostics'] = copy.deepcopy(data)
    return result