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

import copy
import json
import os
import re

from tools.rag.build_chunks import TokenCounter
from tools.retrieval.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', 'clarification', 'named_subjects', 'work_scopes',
          'resolved_query', 'name_resolution')


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

    def whole_prefix(body,end):
        boundaries=[m.end() for m in re.finditer(
            r'\n\s*\n|(?=^[ \t]*(?:\d+[.)]|[-+*])\s+|^#{1,6}\s|^Sekcia:)',body,re.M)
            if m.end()<=end]
        return body[:max(boundaries,default=0)]

    def balanced_sources(sources):
        """Reserve evidence space for every requested person before expanding."""
        suffix='\n[OBSAH ZDROJA OREZANÝ]'
        selected=[{**s,'text':suffix if s.get('text') else ''} for s in sources]
        if not fits(render(selected,truncated=True)):
            raise ValueError('Rozpočet nestačí na metadáta všetkých porovnávaných osôb')
        families={s.get('document_path','').removesuffix('README.md'):s['title'] for s in sources
                  if s.get('title') in base['named_subjects'] and s.get('document_path')}
        groups={}
        for i,s in enumerate(sources):
            owner=next((p for path,p in families.items() if s.get('document_path','').startswith(path)),f'other:{i}')
            groups.setdefault(owner,[]).append(i)

        def used():return render(selected,truncated=True)['context_budget']['used_tokens']

        def expand(index,budget):
            body=sources[index].get('text') or ''
            if not body:return
            ceiling=min(limit,used()+max(0,budget))
            def candidate(text):
                replacement={**selected[index],'text':text+(suffix if len(text)<len(body) else '')}
                return render([*selected[:index],replacement,*selected[index+1:]],truncated=True)
            if candidate(body)['context_budget']['used_tokens']<=ceiling:
                text=body
            else:
                low,high=0,len(body)-1
                while low<high:
                    mid=(low+high+1)//2
                    if candidate(body[:mid])['context_budget']['used_tokens']<=ceiling:low=mid
                    else:high=mid-1
                text=body[:low]
                # Keep a whole paragraph/list item, including wrapped caveats.
                text=whole_prefix(body,low)
            selected[index]={**selected[index],'text':text+(suffix if len(text)<len(body) else '')}

        remaining=limit-used();quota=remaining//max(1,len(groups))
        for indexes in groups.values():
            start=used()
            for pos,index in enumerate(indexes):
                expand(index,max(0,quota-(used()-start))//(len(indexes)-pos))
        # Reuse the space left by short profiles, still sharing it by person.
        remaining=limit-used();quota=remaining//max(1,len(groups))
        for indexes in groups.values():
            start=used()
            for pos,index in enumerate(indexes):
                expand(index,max(0,quota-(used()-start))//(len(indexes)-pos))
        return render(selected,truncated=True)

    # 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) and len(base.get('named_subjects',[]))>1:
            result=balanced_sources(sources)
        elif 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': whole_prefix(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