"""Command-only optimizer prompt adapted from AlgoTune Agent's initial system message."""
from __future__ import annotations

import ast
import textwrap
import warnings

from ...components import PromptBuilder
from ...models import Prompt
from ...records import LoopContext, Selection
from .workspace import numbered, unpack_workspace

SPEND_LIMIT_USD = 1.0
HISTORY_PER_ROLE = 5

_COMMANDS = r"""
Every response must contain a short thought followed by ONE command enclosed by triple
backticks. Put each triple-backtick delimiter on its own line. Do not issue two commands.

Commands:

- edit — replace inclusive lines, create a new file with 0-0, or prepend/replace with 0-N:
```
edit
file: solver.py
lines: 5-7
---
def improved_function():
    pass
---
```
- ls
- view_file <file_name> [start_line] — show at most 100 numbered lines (default 1)
- revert — restore the best-performing measured workspace
- reference <input> — run the task's reference solver
- eval_input <input> — show your output, reference output, stdout, and correctness
- eval — evaluate the current workspace on the task scorer
- delete:
```
delete
file: solver.py
lines: 5-10
```
- profile <filename.py> <input> — profile the current solver
- profile_lines <filename.py> <line[,line|range...]> <input> — focus profiling output

An edit is syntax- and tampering-checked immediately. Every accepted edit/delete is then evaluated
by the task-owned Galapagos AlgoTune scorer; revert reloads the already-measured best snapshot.
Python helpers, setup.py/pyproject.toml, and .pyx/.pxd files form one virtual workspace;
Cython/Pythran/DaCe preparation runs at candidate import time, outside the timed solve call.
Compilation failure returns to the preceding working version.
""".strip()

_HEAD = """SETTING:
You are an autonomous programmer optimizing one AlgoTune task. Every model message costs money;
the result is the best-performing valid code produced at any point, even if a later edit fails.
Keep working until the budget is exhausted.

YOUR TASK:
Optimize the virtual `solver.py` candidate. In this Galapagos port the evaluator entrypoint is:

    run_solver(problem) -> Any

The supplied seed already implements that interface and contains the original task/Solver logic.
Preserve correctness under the task's original validator and make the timed call as fast as
possible. Each instance is limited to 10x its reference runtime. Import/JIT/Cython initialization
is not timed, as in upstream AlgoTune. Do not add an `if __name__ == "__main__"` workflow.
""".strip()


def _number_all(source: str) -> str:
    return "\n".join(f"{index:4d}: {line}" for index, line in enumerate(source.splitlines(), 1))


def _upstream_reference_sections(seed_source: str, evaluator_source: str) -> tuple[str, str]:
    """Extract what AlgoTune itself places in the initial message.

    Upstream inspects the task class and includes the reference ``solve`` method plus the
    ``is_solution`` validator and its helpers, rather than exposing its dataset/evaluation harness.
    The Galapagos seed contains that same verbatim task class, so recover those sections with AST
    without importing task code on the host.
    """
    try:
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", SyntaxWarning)
            tree = ast.parse(seed_source)
        imports = [
            ast.get_source_segment(seed_source, node) or ""
            for node in tree.body
            if isinstance(node, (ast.Import, ast.ImportFrom))
        ]
        candidates: list[ast.ClassDef] = []
        for node in ast.walk(tree):
            if not isinstance(node, ast.ClassDef):
                continue
            names = {
                item.name for item in node.body
                if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef))
            }
            if {"solve", "is_solution"} <= names:
                candidates.append(node)
        task_class = candidates[-1]
        methods = {
            node.name: textwrap.dedent(ast.get_source_segment(seed_source, node) or "")
            for node in task_class.body
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
        }
        imports_text = "\n".join(item for item in imports if item)
        reference = "\n\n".join(item for item in (imports_text, methods["solve"]) if item)
        helpers = [
            source for name, source in methods.items()
            if name not in {"generate_problem", "solve", "is_solution"}
        ]
        validator = "\n\n".join(
            item for item in (imports_text, *helpers, methods["is_solution"]) if item
        )
        return reference, validator
    except (SyntaxError, KeyError, IndexError):
        # A custom AlgoTune-compatible task may not embed the canonical class. Its seed and evaluator
        # remain the only honest sources available, so retain them rather than inventing an interface.
        return seed_source, evaluator_source


class AlgoTunePromptBuilder(PromptBuilder):
    """Render the immutable initial prompt; command results carry subsequent turns."""

    def __init__(self):
        self.task_name = ""
        self.task_context = ""
        self.reference_source = ""
        self.validator_source = ""

    def bind_task(self, task) -> None:
        self.task_name = str(getattr(task, "name", ""))
        self.task_context = str(getattr(task, "context", "") or "")
        seed = str(getattr(task, "initial_program_source", "") or "")
        evaluator = str(getattr(task, "evaluator_source", "") or "")
        self.reference_source, self.validator_source = _upstream_reference_sections(seed, evaluator)

    def build(self, selection: Selection, memory=None, ctx: LoopContext | None = None) -> Prompt:
        parent = selection.parent
        current = unpack_workspace(parent.content) if parent is not None else {"solver.py": ""}
        files = "\n".join(
            f"- {name} ({len(text.splitlines())} lines)"
            for name, text in sorted(current.items())
        )
        solver_view = numbered(current.get("solver.py", ""), 1, 100)
        task_context = self.task_context or (ctx.task_context if ctx else "")
        system = f"""{_HEAD}

{_COMMANDS}

GOALS:
Return the same valid output as the reference and maximize speedup. A wrong instance voids the
entire speed score. Use `reference`, `eval_input`, and profiling to investigate before guessing.

TASK NAME:
{self.task_name}

TASK DESCRIPTION:
{task_context}

REFERENCE CANDIDATE (numbered):
{_number_all(self.reference_source)}

TRUSTED VALIDATOR / SCORER (numbered):
{_number_all(self.validator_source)}
"""
        spent = float(ctx.cost_usd if ctx else 0.0)
        user = f"""Current virtual workspace:
{files}

solver.py (first 100 lines):
{solver_view}

Budget used: ${spent:.4f} / ${SPEND_LIMIT_USD:.4f}; remaining: ${max(0.0, SPEND_LIMIT_USD - spent):.4f}.
Issue exactly one command now."""
        return Prompt(system=system, user=user)