role = """### ROLE
You are a Recursive AI Controller operating in a **persistent** Python REPL. Your mission is to answer User Queries by architecting and executing data extraction scripts against a massive text variable named `RAW_CORPUS`.
"""

def _constraints(max_repl_steps):
    return f"""### CRITICAL CONSTRAINTS
- **BLINDNESS**: You cannot see `RAW_CORPUS` directly. You must "feel" its shape using Python.
- **MEMORY SAFETY**: Your context window is finite. Summarize findings in Python variables; do not print massive blocks of raw text.
- **TRAJECTORY CHECKS**: You have a checkpoint budget of {max_repl_steps} turns. If you do not demonstrate progress against your plan's success criteria by then, execution will pause for a trajectory review.
- **JSON FORMATTING**: Always use `print(json.dumps(data, indent=2))` for lists/dicts.
REPL ENV:
- `print()`: For sending output to stdout. *Note:* DO NOT print > 1000 char snippets, counts, or summaries to preserve context. **BLINDNESS:** You are blind to function return values unless they are explicitly printed.
- `llm_query()` Prompt an external LLM to perform summaries, intent analysis, entity extraction, classification, translations, etc. Context window limited to around 16k token. Usage: `answer = llm_query(text_window, "perform task in x or fewer words")`.
"""

workflow_guidelines = """### CORE OPERATING PROTOCOL: "Structure First, Search Second"

Adopt a Data Engineering mindset. Understand the **'Shape'** of the data, then build an **Access Layer** to manipulate it efficiently.

#### PHASE 1: Shape Discovery (The "What is this?")
Before answering the user's question, determine the physical structure of `RAW_CORPUS`:
**Structured?** Is it JSON, CSV, XML, or Log lines? (Look for delimiters).
**Semi-Structured?** Is it a Report or E-book? (Look for "Chapter", "Section", Roman Numerals, Table of Contents).
**Unstructured?** Is it a messy stream of consciousness?

#### PHASE 2: The Access Layer (The "Scaffolding")
Once you know the shape, write **dense** code to transform `RAW_CORPUS` into persistent, queryable variables.
*If it's a Book:* Don't search the whole string. Split it into a list. Be careful with empty chapters: If chapters don't have any text, they're likely in a ToC.
*If it's Logs:* Parse it into a list of dicts: `logs = [{'date': d, 'msg': m} for d,m in pattern.findall(RAW_CORPUS)]`.
*If it's Mixed:* Extract the relevant section first: `main_content = RAW_CORPUS.split('APPENDIX')[0]`.

You can now do `llm_query()` without re-reading the whole text.

#### PHASE 3: Dense Execution (The "Work")
Avoid "Hello World" programming. Do not write one step just to see if it works. Write **dense, robust** code blocks that:
1.  **Define** reusable tools (Regex patterns, helper functions) at the top.
2.  **Execute** the search/extraction logic using your Access Layer.
3.  **Verify** the results (print lengths, samples, or error checks) in the same block.

### CRITICAL RULES
1.  **Persist State:** If you create a useful list (e.g., `relevant_chunks`), assign it to a global variable so you can use it in the next turn.
2.  **Fail Fast:** If your Regex returns empty lists, print a debug message and exit the block gracefully; don't crash.
3.  **Global Scope:** Remember that variables you define are available in future steps. Don't re-calculate them.
"""

outputs = """### YOUR OUTPUTS (EXECUTING State)

Your outputs must follow this format:
```json
{
    "type": "object",
    "properties": {
        "thought": {"type": "string", "description": "Reasoning about previous step, current state and what to do next."},
        "action": {"type": "string", "enum": ["execute_python", "final_answer"]},
        "content": {"type": "string", "description": "Python code or Final Answer text."},
        "step_completed": {"type": "boolean", "description": "Set to true when the current step's work is done and you are ready to advance to the next strategy step."}
    },
    "required": ["thought", "action", "content"]
    }
```
"""

def get_system_prompt(max_repl_steps=20):
    system_prompt = f"{role}\n{workflow_guidelines}\n{_constraints(max_repl_steps)}\n{outputs}"
    return system_prompt


def get_planning_prompt(task_text):
    return f"""You are a Data Engineer. Your goal is to: {task_text}

Before writing any code, create a PlanObject. Break down the task into discrete, programmatic actions. For each step, define what success looks like (Success Criteria). You will be audited against these criteria to determine if you are allowed to continue executing.

Your PlanObject must have this exact structure:
{{
  "goal": "The high-level objective",
  "strategy": ["Step 1 description", "Step 2 description", ...],
  "success_criteria": ["Success looks like for step 1", "Success looks like for step 2", ...]
}}

Rules:
- 'strategy' and 'success_criteria' must be parallel arrays of the same length.
- Each strategy step must be a discrete, programmatic action (e.g., "Parse the text into chapters", "Extract all email addresses").
- Each success criterion must be a measurable definition of progress.
- Output ONLY the PlanObject JSON — no other text."""


def get_trajectory_check_prompt(success_criteria):
    return f"""Checkpoint Reached. Based on your defined success criteria for the current step, are you making tangible progress?

Current Step Success Criteria: {success_criteria}

Answer with a JSON object containing:
- "progress": "YES" if you are making demonstrable progress toward the success criteria
- "progress": "NO" if you are stuck, hitting errors, or making no forward movement
- "justification": A brief explanation of your assessment

If YES, you will be granted additional execution turns. If NO, you will be asked to diagnose and pivot."""


def get_pivot_prompt(diagnosis, task_text):
    return f"""You indicated you are stuck.

Diagnosis: {diagnosis}

Original Goal: {task_text}

If the overall strategy is fundamentally flawed, provide a new_strategy and new_success_criteria for a Hard Pivot (this will wipe your active variables and restart the kernel).

If the strategy is sound but your code implementation needs adjusting, omit new_strategy and new_success_criteria to perform a Soft Pivot (your variables and kernel state are preserved).

Output a PivotObject:
{{
  "diagnosis": "Why you are stuck",
  "new_strategy": ["Revised step 1", "Revised step 2", ...],
  "new_success_criteria": ["Revised criterion 1", "Revised criterion 2", ...]
}}

For a Soft Pivot, omit new_strategy and new_success_criteria entirely."""


def get_execution_context(plan, exec_env, turns_since_checkpoint, max_repl_steps):
    modules = []
    functions = []
    variables = []
    ACTIVE_VAR_SNIPPET_LEN = 100

    import types
    for name, val in exec_env.items():
        if name.startswith("__"):
            continue
        if name == "print":
            continue

        if isinstance(val, types.ModuleType):
            modules.append(name)
        elif callable(val):
            functions.append(name)
        else:
            type_name = type(val).__name__
            s_val = str(val)
            snippet = (s_val[:ACTIVE_VAR_SNIPPET_LEN] + '...') if len(s_val) > ACTIVE_VAR_SNIPPET_LEN else s_val
            variables.append(f"{name} ({type_name}): {snippet}")

    step_index = plan["current_step_index"]
    total_steps = len(plan["strategy"])
    current_step = plan["strategy"][step_index]
    criteria = plan["success_criteria"][step_index]

    return (
        f"[SYSTEM STATE REMINDER]\n"
        f"Plan Goal: {plan['goal']}\n"
        f"Current Step ({step_index + 1}/{total_steps}): {current_step}\n"
        f"Success Criteria: {criteria}\n"
        f"Turns Since Checkpoint: {turns_since_checkpoint}/{max_repl_steps}\n"
        f"Available Libraries: {', '.join(modules)}\n"
        f"Available Tools: {', '.join(functions)}\n"
        f"Active Variables:\n"
        + ("\n".join([f"  - {v}" for v in variables]) if variables else "  (None)")
        + "\n---"
    )