from __future__ import annotations

import json

from app.modules.rag.explain.models import ExplainPack


class PromptBudgeter:
    def __init__(
        self,
        *,
        max_paths: int = 3,
        max_symbols: int = 25,
        max_excerpts: int = 40,
        max_chars: int = 30000,
    ) -> None:
        self._max_paths = max_paths
        self._max_symbols = max_symbols
        self._max_excerpts = max_excerpts
        self._max_chars = max_chars

    def build_prompt_input(self, question: str, pack: ExplainPack) -> str:
        symbol_ids: list[str] = []
        for path in pack.trace_paths[: self._max_paths]:
            for symbol_id in path.symbol_ids:
                if symbol_id and symbol_id not in symbol_ids and len(symbol_ids) < self._max_symbols:
                    symbol_ids.append(symbol_id)
        excerpts = []
        total_chars = 0
        for excerpt in pack.code_excerpts:
            if symbol_ids and excerpt.symbol_id and excerpt.symbol_id not in symbol_ids:
                continue
            body = excerpt.content.strip()
            remaining = self._max_chars - total_chars
            if remaining <= 0 or len(excerpts) >= self._max_excerpts:
                break
            if len(body) > remaining:
                body = body[:remaining].rstrip() + "...[truncated]"
            excerpts.append(
                {
                    "evidence_id": excerpt.evidence_id,
                    "title": excerpt.title,
                    "path": excerpt.path,
                    "start_line": excerpt.start_line,
                    "end_line": excerpt.end_line,
                    "focus": excerpt.focus,
                    "content": body,
                }
            )
            total_chars += len(body)
        payload = {
            "question": question,
            "intent": pack.intent.model_dump(mode="json"),
            "selected_entrypoints": [item.model_dump(mode="json") for item in pack.selected_entrypoints[:5]],
            "seed_symbols": [item.model_dump(mode="json") for item in pack.seed_symbols[: self._max_symbols]],
            "trace_paths": [path.model_dump(mode="json") for path in pack.trace_paths[: self._max_paths]],
            "evidence_index": {key: value.model_dump(mode="json") for key, value in pack.evidence_index.items()},
            "code_excerpts": excerpts,
            "missing": pack.missing,
            "conflicts": pack.conflicts,
        }
        return json.dumps(payload, ensure_ascii=False, indent=2)