from __future__ import annotations

import json
from typing import Any, Callable, Dict, Optional


TREE_SEARCH_PROMPT_V1 = """You are given a question and the tree structure of a document.
Each node contains node id, title, page range, and summary fields.
Identify all nodes likely to contain the answer.

Question: {query}

Document tree structure:
{tree_json}

Return ONLY valid JSON in this schema:
{{
  "thinking": "brief reasoning about relevant nodes",
  "node_list": ["0001", "0002"]
}}
"""


ANSWER_PROMPT_V1 = """Answer the question based only on the context.
If the context is insufficient, say exactly what is missing.

Question: {query}

Context:
{context}

Return a clear, concise answer grounded in the provided context.
"""


def render_tree_search_prompt(
    query: str,
    tree_view: Dict[str, Any],
    template: Optional[str | Callable[[str, Dict[str, Any], Dict[str, Any]], str]] = None,
    extra: Optional[Dict[str, Any]] = None,
) -> str:
    payload = extra or {}
    if callable(template):
        return template(query, tree_view, payload)

    chosen = template or TREE_SEARCH_PROMPT_V1
    tree_json = json.dumps(tree_view, indent=2, ensure_ascii=False)
    return chosen.format(
        query=query,
        tree_json=tree_json,
        extra_json=json.dumps(payload, ensure_ascii=False),
    )


def render_answer_prompt(
    query: str,
    context_text: str,
    template: Optional[str | Callable[[str, str, Dict[str, Any]], str]] = None,
    extra: Optional[Dict[str, Any]] = None,
) -> str:
    payload = extra or {}
    if callable(template):
        return template(query, context_text, payload)

    chosen = template or ANSWER_PROMPT_V1
    return chosen.format(
        query=query,
        context=context_text,
        extra_json=json.dumps(payload, ensure_ascii=False),
    )
