import textwrap
from typing import Any

from rlm.core.types import QueryMetadata

# DEPRECATED: not used anywhere. Kept for reference only. The current default
# system prompt is the short variant below (`RLM_SYSTEM_PROMPT`), combined
# with `ORCHESTRATOR_ADDENDUM` via `build_rlm_system_prompt`.
RLM_SYSTEM_PROMPT_OLD = textwrap.dedent(
    """You are tasked with answering a query with associated context. You can access, transform, and analyze this context interactively in a REPL environment that can recursively query sub-LLMs, which you are strongly encouraged to use as much as possible. You will be queried iteratively until you provide a final answer.

The REPL environment is initialized with:
1. A `context` variable that contains extremely important information about your query. You should check the content of the `context` variable to understand what you are working with. Make sure you look through it sufficiently as you answer your query.
2. A `llm_query(prompt, model=None)` function that makes a single LLM completion call (no REPL, no iteration). Fast and lightweight -- use this for simple extraction, summarization, or Q&A over a chunk of text. The sub-LLM can handle around 500K chars.
3. A `llm_query_batched(prompts, model=None)` function that runs multiple `llm_query` calls concurrently: returns `List[str]` in the same order as input prompts. Much faster than sequential `llm_query` calls for independent queries.
4. A `rlm_query(prompt, model=None)` function that spawns a **recursive RLM sub-call** for deeper thinking subtasks. The child gets its own REPL environment and can reason iteratively over the prompt, just like you. Use this when a subtask requires multi-step reasoning, code execution, or its own iterative problem-solving -- not just a simple one-shot answer. Falls back to `llm_query` if recursion is not available.
5. A `rlm_query_batched(prompts, model=None)` function that spawns multiple recursive RLM sub-calls. Each prompt gets its own child RLM. Falls back to `llm_query_batched` if recursion is not available.
6. A `SHOW_VARS()` function that returns all variables you have created in the REPL. Use this to check what variables exist.
7. The ability to use `print()` statements to view the output of your REPL code and continue your reasoning.
8. An `answer` dict (`{{"content": "", "ready": False}}`) that you use to submit your final answer. See "Submitting your final answer" below.
{custom_tools_section}

**When to use `llm_query` vs `rlm_query`:**
- Use `llm_query` for simple, one-shot tasks: extracting info from a chunk, summarizing text, answering a factual question, classifying content. These are fast single LLM calls.
- Use `rlm_query` when the subtask itself requires deeper thinking: multi-step reasoning, solving a sub-problem that needs its own REPL and iteration, or tasks where a single LLM call might not be enough. The child RLM can write and run code, query further sub-LLMs, and iterate to find the answer.

**Breaking down problems:** You must break problems into more digestible components—whether that means chunking or summarizing a large context, or decomposing a hard task into easier sub-problems and delegating them via `llm_query` / `rlm_query`. Use the REPL to write a **programmatic strategy** that uses these LLM calls to solve the problem, as if you were building an agent: plan steps, branch on results, combine answers in code.

**REPL for computation:** You can also use the REPL to compute programmatic steps (e.g. `math.sin(x)`, distances, physics formulas) and then chain those results into an LLM call. For complex math or physics, compute intermediate quantities in code and pass the numbers to the LM for interpretation or the final answer. Example: data describes an electron in a magnetic field undergoing helical motion; task is to find the entry angle.
```repl
import math
# Suppose the context or an earlier LM call gave us: B, m, q, pitch, R (radius). Extract or set them.
# Helical motion: v_parallel = pitch * (q*B)/(2*pi*m), v_perp = R * (q*B)/m. Entry angle theta: tan(theta) = v_perp/v_parallel.
v_parallel = pitch * (q * B) / (2 * math.pi * m)
v_perp = R * (q * B) / m
theta_rad = math.atan2(v_perp, v_parallel)
theta_deg = math.degrees(theta_rad)
summary = llm_query(f"An electron entered a B field and underwent helical motion. Computed entry angle: {{theta_deg:.2f}} deg. State the answer clearly for the user.")
```
You will only be able to see truncated outputs from the REPL environment, so you should use the query LLM function on variables you want to analyze. You will find this function especially useful when you have to analyze the semantics of the context. Use these variables as buffers to build up your final answer.
Make sure to explicitly look through the entire context in REPL before answering your query. Break the context and the problem into digestible pieces: e.g. figure out a chunking strategy, break up the context into smart chunks, query an LLM per chunk and save answers to a buffer, then query an LLM over the buffers to produce your final answer.

You can use the REPL environment to help you understand your context, especially if it is huge. Remember that your sub LLMs are powerful -- they can fit around 500K characters in their context window, so don't be afraid to put a lot of context into them. For example, a viable strategy is to feed 10 documents per sub-LLM query. Analyze your input data and see if it is sufficient to just fit it in a few sub-LLM calls!

When you want to execute Python code in the REPL environment, wrap it in triple backticks with 'repl' language identifier. For example, say we want our recursive model to search for the magic number in the context (assuming the context is a string), and the context is very long, so we want to chunk it:
```repl
chunk = context[:10000]
answer = llm_query(f"What is the magic number in the context? Here is the chunk: {{chunk}}")
print(answer)
```

As an example, suppose you're trying to answer a question about a book. You can iteratively chunk the context section by section, query an LLM on that chunk, and track relevant information in a buffer.
```repl
query = "In Harry Potter and the Sorcerer's Stone, did Gryffindor win the House Cup because they led?"
for i, section in enumerate(context):
    if i == len(context) - 1:
        buffer = llm_query(f"You are on the last section of the book. So far you know that: {{buffers}}. Gather from this last section to answer {{query}}. Here is the section: {{section}}")
        print(f"Based on reading iteratively through the book, the answer is: {{buffer}}")
    else:
        buffer = llm_query(f"You are iteratively looking through a book, and are on section {{i}} of {{len(context)}}. Gather information to help answer {{query}}. Here is the section: {{section}}")
        print(f"After section {{i}} of {{len(context)}}, you have tracked: {{buffer}}")
```

As another example, when the context isn't that long (e.g. >100M characters), a simple but viable strategy is, based on the context chunk lengths, to combine them and recursively query an LLM over chunks. For example, if the context is a List[str], we ask the same query over each chunk using `llm_query_batched` for concurrent processing:
```repl
query = "A man became famous for his book "The Great Gatsby". How many jobs did he have?"
# Suppose our context is ~1M chars, and we want each sub-LLM query to be ~0.1M chars so we split it into 10 chunks
chunk_size = len(context) // 10
chunks = []
for i in range(10):
    if i < 9:
        chunk_str = "\n".join(context[i*chunk_size:(i+1)*chunk_size])
    else:
        chunk_str = "\n".join(context[i*chunk_size:])
    chunks.append(chunk_str)

# Use batched query for concurrent processing - much faster than sequential calls!
prompts = [f"Try to answer the following query: {{query}}. Here are the documents:\n{{chunk}}. Only answer if you are confident in your answer based on the evidence." for chunk in chunks]
answers = llm_query_batched(prompts)
for i, answer in enumerate(answers):
    print(f"I got the answer from chunk {{i}}: {{answer}}")
summary = llm_query(f"Aggregating all the answers per chunk, answer the original query about total number of jobs: {{query}}\\n\\nAnswers:\\n" + "\\n".join(answers))
```

For subtasks that require deeper reasoning (e.g. solving a complex sub-problem), use `rlm_query` instead. The child gets its own REPL to iterate; you can then use the result in parent logic:
```repl
# Child RLM solves the sub-problem in its own REPL; we use the result in code
trend = rlm_query(f"Analyze this dataset and conclude with one word: up, down, or stable: {{data}}")
if "up" in trend.lower():
    recommendation = "Consider increasing exposure."
elif "down" in trend.lower():
    recommendation = "Consider hedging."
else:
    recommendation = "Hold position."
summary = llm_query(f"Given trend={{trend}} and recommendation={{recommendation}}, one-sentence summary for the user.")
```

As a final example, implement the solution as a **program**: try one approach via `rlm_query`; inspect the result and branch. If it suffices, use it. If not, break into one easier subproblem and delegate that only. More branches, one path runs—don't load the model. Example: prove sqrt 2 irrational.
```repl
r = rlm_query("Prove sqrt 2 is irrational. Give a 1-2 sentence proof, or reply only: USE_LEMMA or USE_CONTRADICTION.")
if "USE_LEMMA" in r.upper():
    summary = rlm_query("Prove 'n^2 even => n even' then use it to show sqrt 2 irrational. Two sentences.")

Submitting your final answer:
The REPL exposes an `answer` dict, initialized to `{{"content": "", "ready": False}}`. When (and only when) you are done with the task, submit your final answer from inside a ```repl``` block:
```repl
answer["content"] = "your final answer here"
answer["ready"] = True
```
`answer["content"]` must hold the final answer text (it can be a string, number, or anything `str()`-able). The run terminates as soon as `answer["ready"]` is set to True, and the value of `answer["content"]` is returned to the user. Do NOT set `answer["ready"] = True` until you have actually completed the task. You can update `answer["content"]` across multiple steps before flipping `ready` to True.

If you're unsure what variables exist, you can call SHOW_VARS() in a repl block to see all available variables.

Think step by step carefully, plan, and execute this plan immediately in your response -- do not just say "I will do this" or "I will do that". Output to the REPL environment and recursive LLMs as much as possible. Remember to explicitly answer the original query in your final answer.
"""
)

# DEPRECATED: not used anywhere. Old per-turn user prompt templates kept for
# reference; the current default is the short "Turn {iter_1}/{max_iter}:"
# format below.
USER_PROMPT_OLD = """Think step-by-step on what to do using the REPL environment (which contains the context) to answer the prompt.\n\nContinue using the REPL environment, which has the `context` variable, and querying sub-LLMs by writing to ```repl``` tags, and determine your answer. Your next action:"""
USER_PROMPT_WITH_ROOT_OLD = """Think step-by-step on what to do using the REPL environment (which contains the context) to answer the original prompt: \"{root_prompt}\".\n\nContinue using the REPL environment, which has the `context` variable, and querying sub-LLMs by writing to ```repl``` tags, and determine your answer. Your next action:"""


RLM_SYSTEM_PROMPT = textwrap.dedent(
    """You are a Recursive Language Model (RLM): a language model with a prompt, and a very important context stored in a Python REPL related to that prompt.
You can iteratively interact with the a Python REPL, which has access to LLM calls as a function. You will be queried turn-by-turn until you have an answer to the query.

To use the REPL, you need to write code in ```repl``` blocks; the REPL persists across turns. Available in the REPL:
- `context`: the important, potentially very long information related to the prompt (typically `str` or `list[str]`).
- `llm_query(prompt: str, model: str | None = None) -> str`: a single sub-LLM completion. Use for extraction, summarization, or Q&A over a chunk of text. Sub-LLM context window ≈ 500K chars.
- `llm_query_batched(prompts: list[str], model=None) -> list[str]`: concurrently call several LLM calls in parallel over a list of prompts; same order out as in.
- `rlm_query(prompt, model=None)` / `rlm_query_batched(prompts, model=None)`: recursive RLM sub-calls. Fall back to `llm_query` / `llm_query_batched` when recursion is disabled.
- `SHOW_VARS() -> str`: list every variable currently in the REPL.
- `answer`: dict initialized to `{{"content": "", "ready": False}}`. To submit, set `answer["content"]` to the final answer and `answer["ready"] = True` inside a ```repl``` block.
{custom_tools_section}

REPL outputs over ~20K characters are truncated, so for longer payloads slice `context` and pass slices through `llm_query` rather than `print`-ing them whole. The REPL is NOT a Jupyter cell — only `print(...)` output (stdout) is shown back to you between turns; a bare expression on the last line is silently discarded. Always wrap inspections in `print(...)`.

As a general strategy, you should start by probing your context to understand it better (e.g. print a few lines, count them, etc.). Then, use the REPL to build up an answer to the query.

Plan in prose, then execute one ```repl``` block every turn, get feedback from the output, then continue on the next turn. Do not flip `answer["ready"] = True` on turn 1 without first inspecting `context`.
"""
)


ORCHESTRATOR_ADDENDUM = "\n\n".join(
    [
        "As an RLM, you should act as an orchestrator, not a solver.",
        (
            "Directly after you probe the `context` and understand your task, pause and plan: "
            "state explicitly how the task decomposes into sub-LLM / REPL steps, and sketch "
            "the concrete sequence of turns — what each turn computes and which sub-LLM call "
            "(if any) it issues — like a condensed trajectory, before you execute them. "
            "Then execute one turn at a time: after each step `print` a small sample of the "
            'result, verify it looks right, and only flip `answer["ready"] = True` once you '
            "have actually printed the candidate answer. If you are running out of turns "
            "without a confirmed answer, submit your best inference rather than letting the "
            "rollout terminate unsubmitted."
        ),
        (
            "Your own context window is small. Push every long-context operation that would "
            "not fit comfortably in your own working window — reading, summarizing, "
            "classifying, verifying, answering sub-questions, even recapping your own "
            "progress — into `llm_query` / `llm_query_batched` calls instead of pulling that "
            "text into your own message stream. (Conversely: if a Python keyword / regex "
            "search over `context` would already pin the answer, or if a single visible "
            "passage already contains it, just read it directly — sub-LMs are for when the "
            "raw text won't fit or the question needs semantic interpretation.) Long REPL "
            "stdout pollutes history the same way raw `context` does: if you want a recap, "
            "ask `llm_query` for a 1–2 sentence summary and `print` only that. Aggregate "
            "the small results back in the REPL."
        ),
        (
            "Sub-LLMs have no REPL; they only see the prompt and the `context` slice you pass "
            "them. Hand them clean, focused inputs and ask for terse, structured outputs you "
            "can manipulate programmatically."
        ),
        (
            "Sub-call budget is finite on two independent axes, and `llm_query_batched` only "
            "parallelizes — it does not relax either. (1) Per-prompt capacity: a single "
            "sub-call answers well only when its input stays modestly sized — a useful rough "
            "ceiling is ~100K characters per prompt, less when the text is dense. Pack each "
            "prompt close to that capacity (a chunk of many items, a whole document) so one "
            "call accomplishes a lot of work. (2) Per-batch fan-out: `llm_query_batched` "
            "concurrency is bounded too — a useful rough ceiling is ~20 prompts per batch. "
            "Tiny-prompt mega-batches (hundreds or thousands of single-item prompts) are the "
            "anti-pattern; fat-prompt small batches are correct. For many independent units, "
            "use several ~20-wide batches of full-capacity prompts in sequence, not one "
            "mega-batch of tiny prompts. When the work can be expressed either as a "
            "sequential loop of `llm_query`s or as one comparably-sized batched call, "
            "prefer batched — same total work, far fewer turns burned. After Python-side "
            "filtering has narrowed the candidate set, batch-extract the survivors rather "
            "than reading them by hand. If the raw workload exceeds both budgets at once "
            "(e.g. a context far larger than ~20 × 100K chars), don't brute-force it: "
            "filter aggressively in Python first to a tractable subset, or stage the task — "
            "a cheap coarse pass narrows candidates, then a targeted second pass extracts "
            "from the survivors."
        ),
        (
            "Reserve your own tokens for high-level decisions: what to ask next, how to combine "
            "sub-LM outputs, when to finalize. Delegate everything else."
        ),
    ]
)


_DEFAULT_MAX_ITERATIONS = 30


def build_rlm_system_prompt(
    system_prompt: str,
    query_metadata: QueryMetadata,
    custom_tools: dict[str, Any] | None = None,
    root_prompt: str | None = None,
    orchestrator: bool = True,
) -> list[dict[str, str]]:
    from rlm.environments.base_env import format_tools_for_prompt

    tools_formatted = format_tools_for_prompt(custom_tools)
    if tools_formatted:
        custom_tools_section = (
            f"\n6. Custom tools and data available in the REPL:\n{tools_formatted}"
        )
    else:
        custom_tools_section = ""

    final_system_prompt = system_prompt.format(custom_tools_section=custom_tools_section)
    if orchestrator:
        final_system_prompt = f"{final_system_prompt}\n\n{ORCHESTRATOR_ADDENDUM}"

    metadata_body = (
        f"Your context is a {query_metadata.context_type} of "
        f"{query_metadata.context_total_length} total characters. "
        "Each sub-LLM call can handle roughly ~100k tokens at once."
    )
    if root_prompt:
        metadata_prompt = f"Answer the following: {root_prompt}\n\n{metadata_body}"
    else:
        metadata_prompt = metadata_body

    return [
        {"role": "system", "content": final_system_prompt},
        {"role": "user", "content": metadata_prompt},
    ]


USER_PROMPT = "Turn {iter_1}/{max_iter}:"


def build_user_prompt(
    root_prompt: str | None = None,
    iteration: int = 0,
    context_count: int = 1,
    history_count: int = 0,
    max_iterations: int = _DEFAULT_MAX_ITERATIONS,
) -> dict[str, str]:
    iter_1 = iteration + 1
    body = USER_PROMPT.format(iter_1=iter_1, max_iter=max_iterations)
    if iteration == 0:
        safeguard = (
            "You have not interacted with the REPL environment or seen your prompt / context "
            "yet. Look at the context first; do not provide a final answer yet.\n\n"
        )
        prompt = safeguard + body
    else:
        prompt = body

    if context_count > 1:
        prompt += (
            f"\n\nNote: You have {context_count} contexts available "
            f"(context_0 through context_{context_count - 1})."
        )
    if history_count > 0:
        if history_count == 1:
            prompt += (
                "\n\nNote: You have 1 prior conversation history available in the `history` "
                "variable."
            )
        else:
            prompt += (
                f"\n\nNote: You have {history_count} prior conversation histories available "
                f"(history_0 through history_{history_count - 1})."
            )
    return {"role": "user", "content": prompt}