"""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