"""Prompt construction + code extraction for the OpenSleuth agent."""

from __future__ import annotations

import re
from typing import Iterable

SYSTEM_PROMPT = (
    "You are an algorithmic detective. You are given the public signature of a hidden "
    "Python function plus several (input, output) examples observed by probing it. "
    "Your job is to write a Python function that *exactly* reproduces the hidden "
    "function's behavior on all valid inputs. Match its return values AND its "
    "exception types on invalid inputs. Keep your implementation as simple and clean "
    "as possible (it is penalised for being needlessly branchy). Return ONLY the "
    "function definition wrapped in a single ```python ... ``` code block."
)


def build_prompt(target_name: str, signature: str, probes: Iterable[tuple[str, str, bool]]) -> str:
    """Build the user-side prompt.

    `probes` is an iterable of `(input_repr, output_repr, is_error)` tuples,
    typically pre-sampled by the dataset builder.
    """
    lines = [
        f"## Hidden function: {target_name}",
        "",
        f"### Public signature & docstring",
        signature.strip() or "(no signature provided)",
        "",
        "### Observed probes",
    ]
    probe_list = list(probes)
    if not probe_list:
        lines.append("(none)")
    else:
        for inp, out, is_err in probe_list:
            tag = "raises" if is_err else "returns"
            lines.append(f"- input={inp}  ->  {tag} {out}")
    lines += [
        "",
        "### Task",
        f"Write a Python function named `{target_name}` that reproduces the hidden "
        "function's behaviour. Return ONLY the function definition in a single "
        "```python ... ``` code block. Do not add explanations.",
    ]
    return "\n".join(lines)


_CODE_RE = re.compile(r"```(?:python)?\s*(.*?)```", re.DOTALL | re.IGNORECASE)


def extract_code(completion: str) -> str:
    """Pull the python source from a model completion. If no fenced block is
    present we fall back to the whole completion (the verifier will then judge
    it on its own)."""
    m = _CODE_RE.search(completion)
    if m:
        return m.group(1).strip()
    return completion.strip()