#!/usr/bin/env python3
"""Build an audited synthetic tool-use SFT corpus unrelated to Project Euler."""
from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any

TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "write",
            "description": "Write complete content to a file.",
            "parameters": {
                "type": "object",
                "properties": {"path": {"type": "string"}, "content": {"type": "string"}},
                "required": ["path", "content"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "read",
            "description": "Read a text file.",
            "parameters": {
                "type": "object",
                "properties": {"path": {"type": "string"}},
                "required": ["path"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "edit",
            "description": "Replace exact text in a file.",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {"type": "string"},
                    "edits": {
                        "type": "array",
                        "items": {
                            "type": "object",
                            "properties": {
                                "oldText": {"type": "string"},
                                "newText": {"type": "string"},
                            },
                            "required": ["oldText", "newText"],
                        },
                    },
                },
                "required": ["path", "edits"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "bash",
            "description": "Execute a shell command in the workspace.",
            "parameters": {
                "type": "object",
                "properties": {"command": {"type": "string"}},
                "required": ["command"],
            },
        },
    },
]

SYSTEM_PROMPT = (
    "You are a careful coding agent. Use the provided tools to create and inspect files, "
    "execute the assigned program, recover from errors when necessary, and report only "
    "results verified from tool output. Never claim execution without a successful bash result."
)


def tool_call(call_id: str, name: str, arguments: dict[str, Any]) -> dict[str, Any]:
    return {
        "role": "assistant",
        "content": "",
        "tool_calls": [
            {
                "id": call_id,
                "type": "function",
                "function": {"name": name, "arguments": json.dumps(arguments, separators=(",", ":"))},
            }
        ],
    }


def tool_result(call_id: str, name: str, content: str) -> dict[str, Any]:
    return {"role": "tool", "tool_call_id": call_id, "name": name, "content": content}


def fibonacci(index: int) -> int:
    previous, current = 0, 1
    for _ in range(index):
        previous, current = current, previous + current
    return previous


def language_case(language: str, index: int) -> tuple[str, str, str, str, str]:
    limit = 7 + index
    operation = index % 5
    if operation == 0:
        description = f"the sum of squares from 1 through {limit}"
        expression = f"sum(value^2 for value in 1:{limit})"
        python_expression = f"sum(value * value for value in range(1, {limit + 1}))"
        javascript_term = "value * value"
        expected = sum(value * value for value in range(1, limit + 1))
    elif operation == 1:
        description = f"the sum of cubes from 1 through {limit}"
        expression = f"sum(value^3 for value in 1:{limit})"
        python_expression = f"sum(value ** 3 for value in range(1, {limit + 1}))"
        javascript_term = "value ** 3"
        expected = sum(value**3 for value in range(1, limit + 1))
    elif operation == 2:
        description = f"the triangular number for {limit}"
        expression = f"sum(1:{limit})"
        python_expression = f"sum(range(1, {limit + 1}))"
        javascript_term = "value"
        expected = sum(range(1, limit + 1))
    elif operation == 3:
        description = f"Fibonacci number {limit} with F(0)=0 and F(1)=1"
        expression = f"(function fib(n)\n    a, b = 0, 1\n    for _ in 1:n\n        a, b = b, a + b\n    end\n    a\nend)({limit})"
        python_expression = (
            f"(lambda n: (lambda values: values[n])([0, 1] + [0] * max(0, n - 1)))({limit})"
        )
        javascript_term = ""
        expected = fibonacci(limit)
    else:
        description = f"the weighted checksum sum(i * (i + 3)) for i from 1 through {limit}"
        expression = f"sum(value * (value + 3) for value in 1:{limit})"
        python_expression = f"sum(value * (value + 3) for value in range(1, {limit + 1}))"
        javascript_term = "value * (value + 3)"
        expected = sum(value * (value + 3) for value in range(1, limit + 1))

    stem = f"verified_{language}_{index}"
    if language == "python":
        path = f"{stem}.py"
        if operation == 3:
            code = (
                "def fib(n):\n"
                "    previous, current = 0, 1\n"
                "    for _ in range(n):\n"
                "        previous, current = current, previous + current\n"
                "    return previous\n\n"
                f"print(fib({limit}))\n"
            )
        else:
            code = f"print({python_expression})\n"
        command = f"python3 {path}"
    elif language == "javascript":
        path = f"{stem}.js"
        if operation == 3:
            code = (
                "function fib(n) {\n  let previous = 0, current = 1;\n"
                "  for (let i = 0; i < n; i++) [previous, current] = [current, previous + current];\n"
                "  return previous;\n}\n"
                f"console.log(fib({limit}));\n"
            )
        else:
            code = (
                f"let total = 0;\nfor (let value = 1; value <= {limit}; value++) total += {javascript_term};\n"
                "console.log(total);\n"
            )
        command = f"node {path}"
    elif language == "julia":
        path = f"{stem}.jl"
        code = f"println({expression})\n"
        command = f"julia {path}"
    else:
        raise ValueError(f"unsupported language: {language}")
    user = (
        f"Create and execute a {language} program that computes {description}. "
        "Report the result only after execution and finish with the verified value in \\boxed{value}."
    )
    return user, path, code, command, str(expected)


def success_example(language: str, index: int) -> dict[str, Any]:
    user, path, code, command, expected = language_case(language, index)
    write_id = f"call_{language}_{index}_write"
    run_id = f"call_{language}_{index}_run"
    return {
        "metadata": {"kind": "success", "language": language, "index": index},
        "system_prompt": SYSTEM_PROMPT,
        "tools": TOOLS,
        "messages": [
            {"role": "user", "content": user},
            tool_call(write_id, "write", {"path": path, "content": code}),
            tool_result(write_id, "write", f"Successfully wrote {len(code)} bytes to {path}"),
            tool_call(run_id, "bash", {"command": command}),
            tool_result(run_id, "bash", f"{expected}\n"),
            {
                "role": "assistant",
                "content": f"The program executed successfully and printed `{expected}`. The verified result is \\boxed{{{expected}}}.",
            },
        ],
    }


def recovery_example(language: str, index: int) -> dict[str, Any]:
    user, path, code, command, expected = language_case(language, index)
    marker = code.splitlines()[-1]
    if language == "python":
        broken_marker = "print(missing_total)"
        failure = "NameError: name 'missing_total' is not defined"
    elif language == "javascript":
        broken_marker = "console.log(missingTotal);"
        failure = "ReferenceError: missingTotal is not defined"
    else:
        broken_marker = "println(missing_total)"
        failure = "ERROR: LoadError: UndefVarError: `missing_total` not defined"
    broken_code = code.replace(marker, broken_marker)
    prefix = f"call_{language}_{index}"
    return {
        "metadata": {"kind": "recovery", "language": language, "index": index},
        "system_prompt": SYSTEM_PROMPT,
        "tools": TOOLS,
        "messages": [
            {"role": "user", "content": user},
            tool_call(f"{prefix}_write", "write", {"path": path, "content": broken_code}),
            tool_result(f"{prefix}_write", "write", f"Successfully wrote {len(broken_code)} bytes to {path}"),
            tool_call(f"{prefix}_run1", "bash", {"command": command}),
            tool_result(f"{prefix}_run1", "bash", failure),
            {
                "role": "assistant",
                "content": "The run failed because the final print references an undefined name. I will inspect and correct that exact line.",
            },
            tool_call(f"{prefix}_read", "read", {"path": path}),
            tool_result(f"{prefix}_read", "read", broken_code),
            tool_call(
                f"{prefix}_edit",
                "edit",
                {"path": path, "edits": [{"oldText": broken_marker, "newText": marker}]},
            ),
            tool_result(f"{prefix}_edit", "edit", f"Updated {path}"),
            tool_call(f"{prefix}_run2", "bash", {"command": command}),
            tool_result(f"{prefix}_run2", "bash", f"{expected}\n"),
            {
                "role": "assistant",
                "content": f"After correcting the undefined name, the program executed successfully and printed `{expected}`. The verified result is \\boxed{{{expected}}}.",
            },
        ],
    }


def build_examples() -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    for language in ("python", "javascript", "julia"):
        rows.extend(success_example(language, index) for index in range(1, 13))
        rows.extend(recovery_example(language, index) for index in range(13, 16))
    return rows


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path, default=Path("data/tool-bootstrap.jsonl"))
    args = parser.parse_args()
    rows = build_examples()
    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row, separators=(",", ":")) + "\n")
    counts: dict[str, int] = {}
    for row in rows:
        key = f"{row['metadata']['language']}/{row['metadata']['kind']}"
        counts[key] = counts.get(key, 0) + 1
    print(json.dumps({"output": str(args.output), "rows": len(rows), "counts": counts}, indent=2))


if __name__ == "__main__":
    main()