"""Prompt builders for procedure creation and step execution.

These helpers produce **LLM prompts** used by the GA scaffold and your runners:

- :func:`create_direct_prompt` — a simple “ask the question directly” prompt
  (handy for baselines and ARC-style multiple choice).
- :func:`create_procedure_prompt` — asks the model to emit a *global-state*
  procedure JSON that validates your Pydantic schema.
- :func:`create_execution_prompt` — runs a single step by showing inputs,
  the step action, and the required outputs, while reminding the model to
  return **strict JSON** conforming to a supplied schema.

All functions are pure string builders and have no network side-effects.
They are safe to use in tests and with any LLM backend.

Example
-------
>>> from evoproc_procedures.prompts import create_procedure_prompt
>>> p = create_procedure_prompt("Natalia sold clips to 48 friends in April...")
>>> isinstance(p, str)
True
"""

from __future__ import annotations

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

# Import from your plugin package, not from "src."
from evoproc_procedures.models import Procedure

__all__ = [
    "create_direct_prompt",
    "create_procedure_prompt",
    "create_execution_prompt",
]


def create_direct_prompt(item: str) -> str:
    """Build a direct Q→A prompt (baseline / ARC-style).

    Args:
        item: The task text (e.g., question stem or problem statement).

    Returns:
        A concise prompt string suitable for a direct answer call.
    """
    # Keep it minimal to reduce unintended reasoning scaffolds in baselines.
    return f"Solve this problem: {item}"


def create_procedure_prompt(item: str, example_prompt: Optional[str] = None) -> str:
    """Ask the LLM to synthesize a *global-state* procedure JSON.

    The model is instructed to return **exactly one** JSON object that validates
    the Pydantic schema produced by :class:`uhj_procedures.models.Procedure`.

    Args:
        item: Natural-language task to decompose.
        example_prompt: Optional extra hint or in-context example description
            (kept as a single string you can preformat upstream).

    Returns:
        A prompt string suitable for a structured generation call.

    Notes:
        - The schema is injected verbatim via ``Procedure.model_json_schema()``.
        - Constraints match your GA scaffold: step 1 uses ``problem_text`` only,
          later steps may read any earlier variables (global state), and the
          final step emits ``final_answer`` **as a description only**.
    """
    schema_json = json.dumps(Procedure.model_json_schema(), ensure_ascii=False)
    example = f"\n# Example (hint)\n{example_prompt}\n" if example_prompt else ""

    # No fenced code blocks; some providers echo fences literally.
    return(
            "Create a short, executable procedure to solve the problem, but do NOT solve it.\n\n"
            f"# Task\n{item}\n"
            f"{example}\n\n"
            "# Output Contract\n"
            "Return exactly one JSON object that validates against this schema (verbatim):\n"
            f"{schema_json}\n\n"
            "# Global IO Constraints\n"
            "Step 1 inputs: exactly [\"problem_text\"].\n"
            "Inputs resolvable: every step input must come from problem_text or from earlier step outputs.\n"
            "Variable names: snake_case; consistent across steps.\n"
            "No arithmetic evaluation: do not compute numeric results in the procedure.\n\n"
            "# Required structure (must follow)\n"
            "1. Step 1 MUST extract a COMPLETE quantity inventory from problem_text:\n"
                "Create one variable per numeric fact (include units/type in description).\n"
                "Coverage rule: every number in the text must be represented by some Step 1 output variable (or explicitly marked irrelevant in a variable like ignored_facts with brief reasons).\n"
            "2. Next steps MUST define the symbolic relationships needed to compute the unknown:\n"
                "Any compute step MUST include an explicit symbolic expression using variable names (e.g., x = a - b - (c + d)).\n"
                "Do not evaluate.\n"
                "Final step outputs exactly [\"final_answer\"] (definition only, no value).\n\n"
            "# Procedure length\n"
            "Keep it short: target 3–6 steps. Combine related arithmetic into one step as long as you only introduce one new output variable per step.\n"
            "Return ONLY the JSON object.\n"
        )


def create_execution_prompt(
    visible_inputs: Dict[str, Any],
    action: str,
    schema: Dict[str, Any],
    expected_outputs: Iterable[str],
    output_descriptions: Optional[Dict[str, str]] = None,
    *,
    is_final: bool = False,
) -> str:
    """Build a prompt to execute a **single step** with strict JSON output.

    Args:
        visible_inputs: The subset of global state the step may read (already extracted variables).
        action: The step's imperative instruction (e.g., \"extract the two numbers a and b\").
        schema: JSON Schema that the **output object** must validate against
            (for regular steps use your step-output schema; for the last step
            you can pass your answer schema, e.g. GSM/ARC).
        expected_outputs: Names of keys the model must return in the JSON object.
        output_descriptions: Optional mapping of key → human description (shown to the model).
        is_final: If ``True``, annotate that these outputs constitute the final answer.

    Returns:
        A prompt string that instructs the model to return only a JSON object
        matching ``schema`` and containing exactly ``expected_outputs``.
    """
    output_lines = []
    output_descriptions = output_descriptions or {}
    for name in expected_outputs:
        desc = output_descriptions.get(name, "")
        output_lines.append(f"- {name}: {desc}".rstrip())

    outputs_block = "\n".join(output_lines) if output_lines else "(see schema)"
    inputs_block = json.dumps(visible_inputs, ensure_ascii=False, indent=2)
    schema_block = json.dumps(schema, ensure_ascii=False, indent=2)

    final_note = " (final_answer)" if is_final else ""
    return (
        f"{action}\n"
        f"# Inputs (JSON)\n{inputs_block}\n"
        "# Required Outputs\n"
        f"Return a JSON object with exactly these keys{final_note}:\n"
        f"{outputs_block}\n\n"
        "# Format\n"
        "- Return ONLY a JSON object that conforms to the provided schema.\n"
        "- Do NOT include extra keys.\n"
        "- Do NOT include commentary.\n\n"
        "# Schema (verbatim)\n"
        f"{schema_block}"
    )
