from __future__ import annotations

import asyncio
import copy
import json
from typing import Any

from agent_platform.auth import UserIdentity
from agent_platform.models import get_model_spec
from agent_platform.runtime.loop import (
    AgentLoop,
    is_substantive_execution,
    normalize_write_file_content,
    python_source_error,
    select_tool_specs,
    tool_event_details,
)
from agent_platform.runtime.tools import TOOL_METADATA
from agent_platform.store import RuntimeStore


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


def test_write_file_normalizes_double_escaped_code_layout() -> None:
    escaped = r'import time\n\n# benchmark\ndef main():\n\tprint("keep\ninside")\n\nmain()\n'
    assert normalize_write_file_content("benchmark.py", escaped) == (
        'import time\n\n# benchmark\ndef main():\n\tprint("keep\\ninside")\n\nmain()\n'
    )


def test_write_file_normalizes_mixed_code_layout_without_touching_strings() -> None:
    escaped = 'import time\nvalue = "keep\\\\ninside"\\nprint(value)\\n'
    assert normalize_write_file_content("benchmark.py", escaped) == (
        'import time\nvalue = "keep\\\\ninside"\nprint(value)\n'
    )


def test_write_file_consumes_multiple_layout_escape_layers() -> None:
    escaped = r"import time\\n\\ndef main():\\n    print('ok')\\n"
    assert normalize_write_file_content("benchmark.py", escaped) == ("import time\n\ndef main():\n    print('ok')\n")


def test_write_file_normalizes_double_escaped_plain_text() -> None:
    assert normalize_write_file_content("report.md", r"# Report\n\nMeasured: 1.2 s\n") == (
        "# Report\n\nMeasured: 1.2 s\n"
    )


def test_invalid_complete_python_source_is_detected_before_write() -> None:
    error = python_source_error({"path": "benchmark.py", "content": "# Report\n- O(n²)\n"})
    assert error is not None
    assert "invalid character" in error
    assert python_source_error({"path": "benchmark.py", "content": "print('ok')\n"}) is None


def test_read_only_shell_commands_do_not_satisfy_execution_evidence() -> None:
    assert not is_substantive_execution("exec", {"command": "cat report.txt"})
    assert not is_substantive_execution("exec", {"command": "sed -n '1,20p' script.py"})
    assert is_substantive_execution("exec", {"command": "python3 script.py"})
    assert is_substantive_execution("start_process", {"command": "python3 server.py"})


def test_simple_script_task_receives_a_focused_tool_menu() -> None:
    specs = [metadata.openai_spec() for metadata in TOOL_METADATA.values()]
    selected = select_tool_specs(specs, "写脚本对比排序算法，给一个报告", depth=0)
    names = {spec["function"]["name"] for spec in selected}
    assert names == {
        "workspace_status",
        "list_files",
        "read_file",
        "search_files",
        "write_file",
        "apply_patch",
        "exec",
        "update_plan",
    }


def test_complex_requests_enable_only_relevant_optional_tools() -> None:
    specs = [metadata.openai_spec() for metadata in TOOL_METADATA.values()]
    selected = select_tool_specs(
        specs,
        "重构 Git 仓库里的 web server，并记住这个偏好，分工并行调研",
        depth=0,
    )
    names = {spec["function"]["name"] for spec in selected}
    assert {"delegate_task", "git_status", "git_diff", "start_process", "poll_process"} <= names
    assert {"remember", "recall_memory", "forget_memory"} <= names


class ScriptedProvider:
    def __init__(self, responses: list[dict[str, Any]]) -> None:
        self.responses = responses
        self.requests = []

    async def complete(self, **kwargs):
        self.requests.append(copy.deepcopy(kwargs))
        return self.responses.pop(0)


class ParallelRegistry:
    def __init__(self) -> None:
        self.started = 0
        self.both_started = asyncio.Event()

    def specs(self, *, read_only=False, allow_delegate=True):
        return [TOOL_METADATA["read_file"].openai_spec()]

    async def execute(self, name, arguments, context):
        self.started += 1
        if self.started == 2:
            self.both_started.set()
        await asyncio.wait_for(self.both_started.wait(), timeout=1)
        return {"ok": True, "path": arguments["path"]}


class CountingRegistry:
    def __init__(self) -> None:
        self.calls = 0

    def specs(self, *, read_only=False, allow_delegate=True):
        return [TOOL_METADATA["read_file"].openai_spec()]

    async def execute(self, name, arguments, context):
        self.calls += 1
        return {"ok": True, "path": arguments["path"]}


async def test_tool_call_batch_is_bounded(settings) -> None:
    crowded = {
        "choices": [
            {
                "message": {
                    "role": "assistant",
                    "content": None,
                    "tool_calls": [tool_call(str(index), "read_file", {"path": f"{index}.txt"}) for index in range(20)],
                },
                "finish_reason": "tool_calls",
            }
        ]
    }
    final = {"choices": [{"message": {"role": "assistant", "content": "done"}, "finish_reason": "stop"}]}
    store = RuntimeStore(settings.database_url)
    await store.initialize()
    registry = CountingRegistry()
    loop = AgentLoop(ScriptedProvider([crowded, final]), registry, store, max_tool_output_chars=10_000)
    events = []

    async def callback(event_type, payload):
        events.append((event_type, payload))

    try:
        answer = await loop.run(
            spec=get_model_spec("work-medium"),
            messages=[{"role": "user", "content": "检查当前状态"}],
            identity=UserIdentity("u1", "u1@example.test", "U1", "user"),
            raw_user_jwt="jwt",
            chat_id="bounded-batch",
            callback=callback,
        )
    finally:
        await store.close()

    assert answer == "done"
    assert registry.calls == 8
    limited = next(payload for event_type, payload in events if event_type == "tool.batch_limited")
    assert limited == {"requested": 20, "accepted": 8, "iteration": 1, "depth": 0}


async def test_read_only_tool_calls_run_in_parallel(settings) -> None:
    first = {
        "choices": [
            {
                "message": {
                    "role": "assistant",
                    "content": None,
                    "tool_calls": [
                        tool_call("1", "read_file", {"path": "a"}),
                        tool_call("2", "read_file", {"path": "b"}),
                    ],
                },
                "finish_reason": "tool_calls",
            }
        ]
    }
    final = {"choices": [{"message": {"role": "assistant", "content": "done"}, "finish_reason": "stop"}]}
    store = RuntimeStore(settings.database_url)
    await store.initialize()
    registry = ParallelRegistry()
    loop = AgentLoop(ScriptedProvider([first, final]), registry, store, max_tool_output_chars=10_000)
    events = []

    async def callback(event_type, payload):
        events.append(event_type)

    try:
        answer = await loop.run(
            spec=get_model_spec("work-light"),
            messages=[{"role": "user", "content": "inspect both"}],
            identity=UserIdentity("u1", "", "", "user"),
            raw_user_jwt="jwt",
            chat_id="c1",
            callback=callback,
        )
        assert answer == "done"
        assert registry.started == 2
        assert events.count("tool.completed") == 2
    finally:
        await store.close()


class SerialRegistry:
    def __init__(self) -> None:
        self.active = 0
        self.max_active = 0

    def specs(self, *, read_only=False, allow_delegate=True):
        return [TOOL_METADATA["write_file"].openai_spec()]

    async def execute(self, name, arguments, context):
        self.active += 1
        self.max_active = max(self.max_active, self.active)
        await asyncio.sleep(0.02)
        self.active -= 1
        return {"ok": True}


async def test_mutating_tool_calls_are_serialized(settings) -> None:
    first = {
        "choices": [
            {
                "message": {
                    "role": "assistant",
                    "content": None,
                    "tool_calls": [
                        tool_call("1", "write_file", {"path": "a", "content": "a"}),
                        tool_call("2", "write_file", {"path": "b", "content": "b"}),
                    ],
                },
                "finish_reason": "tool_calls",
            }
        ]
    }
    verify = {
        "choices": [
            {
                "message": {
                    "role": "assistant",
                    "content": None,
                    "tool_calls": [tool_call("3", "read_file", {"path": "a"})],
                },
                "finish_reason": "tool_calls",
            }
        ]
    }
    final = {"choices": [{"message": {"role": "assistant", "content": "done"}, "finish_reason": "stop"}]}
    store = RuntimeStore(settings.database_url)
    await store.initialize()
    registry = SerialRegistry()
    loop = AgentLoop(ScriptedProvider([first, verify, final]), registry, store, max_tool_output_chars=10_000)

    async def callback(event_type, payload):
        return None

    try:
        await loop.run(
            spec=get_model_spec("work-light"),
            messages=[{"role": "user", "content": "write both"}],
            identity=UserIdentity("u1", "", "", "user"),
            raw_user_jwt="jwt",
            chat_id="c1",
            callback=callback,
        )
        assert registry.max_active == 1
    finally:
        await store.close()


class RecordingRegistry:
    def __init__(self) -> None:
        self.calls = []

    def specs(self, *, read_only=False, allow_delegate=True):
        return [metadata.openai_spec() for metadata in TOOL_METADATA.values()]

    async def execute(self, name, arguments, context):
        self.calls.append((name, arguments))
        return {"ok": True, "output": f"{name} complete"}


async def test_artifact_completion_is_rejected_until_written_and_verified(settings) -> None:
    premature = {"choices": [{"message": {"role": "assistant", "content": "报告已经完成"}, "finish_reason": "stop"}]}
    write = {
        "choices": [
            {
                "message": {
                    "role": "assistant",
                    "tool_calls": [tool_call("1", "write_file", {"path": "report.md", "content": "ok"})],
                },
                "finish_reason": "tool_calls",
            }
        ]
    }
    verify = {
        "choices": [
            {
                "message": {
                    "role": "assistant",
                    "tool_calls": [tool_call("2", "read_file", {"path": "report.md"})],
                },
                "finish_reason": "tool_calls",
            }
        ]
    }
    final = {"choices": [{"message": {"role": "assistant", "content": "已完成 report.md"}, "finish_reason": "stop"}]}
    provider = ScriptedProvider([premature, write, verify, final])
    registry = RecordingRegistry()
    store = RuntimeStore(settings.database_url)
    await store.initialize()
    events = []

    async def callback(event_type, payload):
        events.append(event_type)

    try:
        answer = await AgentLoop(provider, registry, store, max_tool_output_chars=10_000).run(
            spec=get_model_spec("work-light"),
            messages=[{"role": "user", "content": "写一个报告"}],
            identity=UserIdentity("u1", "", "", "user"),
            raw_user_jwt="jwt",
            chat_id="c1",
            callback=callback,
        )
        assert answer == "已完成 report.md"
        assert [name for name, _ in registry.calls] == ["write_file", "read_file"]
        assert "completion.rejected" in events
        assert provider.requests[1]["messages"][-1]["role"] == "user"
        assert "MUST call" in provider.requests[1]["messages"][-1]["content"]
    finally:
        await store.close()


class FailingExecutionRegistry(RecordingRegistry):
    def __init__(self) -> None:
        super().__init__()
        self.exec_calls = 0

    async def execute(self, name, arguments, context):
        self.calls.append((name, arguments))
        if name == "exec":
            self.exec_calls += 1
            if self.exec_calls == 1:
                return {"ok": False, "error": "SyntaxError: invalid source"}
        return {"ok": True, "output": f"{name} complete"}


async def test_failed_execution_must_be_repaired_before_completion(settings) -> None:
    write = {
        "choices": [
            {
                "message": {
                    "role": "assistant",
                    "tool_calls": [tool_call("1", "write_file", {"path": "sort.py", "content": "broken"})],
                },
                "finish_reason": "tool_calls",
            }
        ]
    }
    read = {
        "choices": [
            {
                "message": {
                    "role": "assistant",
                    "tool_calls": [tool_call("2", "read_file", {"path": "sort.py"})],
                },
                "finish_reason": "tool_calls",
            }
        ]
    }
    failed_exec = {
        "choices": [
            {
                "message": {
                    "role": "assistant",
                    "tool_calls": [tool_call("3", "exec", {"command": "python3 sort.py"})],
                },
                "finish_reason": "tool_calls",
            }
        ]
    }
    premature = {"choices": [{"message": {"role": "assistant", "content": "已完成"}, "finish_reason": "stop"}]}
    repair = {
        "choices": [
            {
                "message": {
                    "role": "assistant",
                    "tool_calls": [tool_call("4", "write_file", {"path": "sort.py", "content": "fixed"})],
                },
                "finish_reason": "tool_calls",
            }
        ]
    }
    repaired_exec = {
        "choices": [
            {
                "message": {
                    "role": "assistant",
                    "tool_calls": [tool_call("5", "exec", {"command": "python3 sort.py"})],
                },
                "finish_reason": "tool_calls",
            }
        ]
    }
    final = {
        "choices": [{"message": {"role": "assistant", "content": "脚本运行和报告验证完成"}, "finish_reason": "stop"}]
    }
    provider = ScriptedProvider([write, read, failed_exec, premature, repair, repaired_exec, final])
    registry = FailingExecutionRegistry()
    store = RuntimeStore(settings.database_url)
    await store.initialize()
    events = []

    async def callback(event_type, payload):
        events.append(event_type)

    try:
        answer = await AgentLoop(provider, registry, store, max_tool_output_chars=10_000).run(
            spec=get_model_spec("work-light"),
            messages=[{"role": "user", "content": "写脚本对比排序算法并生成报告"}],
            identity=UserIdentity("u1", "", "", "user"),
            raw_user_jwt="jwt",
            chat_id="c1",
            callback=callback,
        )
        assert answer == "脚本运行和报告验证完成"
        assert registry.exec_calls == 2
        assert events.count("completion.rejected") == 1
        direct_recovery = provider.requests[3]["messages"][-1]
        assert direct_recovery["role"] == "user"
        assert "Execution recovery is required" in direct_recovery["content"]
        recovery = provider.requests[4]["messages"][-1]
        assert recovery["role"] == "user"
        assert "SyntaxError" in recovery["content"]
    finally:
        await store.close()


async def test_unchanged_failed_command_is_blocked_until_a_repair(settings) -> None:
    write = {
        "choices": [
            {
                "message": {
                    "role": "assistant",
                    "tool_calls": [tool_call("1", "write_file", {"path": "sort.py", "content": "broken"})],
                },
                "finish_reason": "tool_calls",
            }
        ]
    }
    failed_exec = {
        "choices": [
            {
                "message": {
                    "role": "assistant",
                    "tool_calls": [tool_call("2", "exec", {"command": "python3 sort.py"})],
                },
                "finish_reason": "tool_calls",
            }
        ]
    }
    repeated_exec = {
        "choices": [
            {
                "message": {
                    "role": "assistant",
                    "tool_calls": [tool_call("3", "exec", {"command": "python3 sort.py"})],
                },
                "finish_reason": "tool_calls",
            }
        ]
    }
    repair = {
        "choices": [
            {
                "message": {
                    "role": "assistant",
                    "tool_calls": [tool_call("4", "write_file", {"path": "sort.py", "content": "fixed"})],
                },
                "finish_reason": "tool_calls",
            }
        ]
    }
    successful_exec = {
        "choices": [
            {
                "message": {
                    "role": "assistant",
                    "tool_calls": [tool_call("5", "exec", {"command": "python3 sort.py"})],
                },
                "finish_reason": "tool_calls",
            }
        ]
    }
    final = {"choices": [{"message": {"role": "assistant", "content": "验证完成"}, "finish_reason": "stop"}]}
    provider = ScriptedProvider([write, failed_exec, repeated_exec, repair, successful_exec, final])
    registry = FailingExecutionRegistry()
    store = RuntimeStore(settings.database_url)
    await store.initialize()

    async def callback(event_type, payload):
        return None

    try:
        answer = await AgentLoop(provider, registry, store, max_tool_output_chars=10_000).run(
            spec=get_model_spec("work-light"),
            messages=[{"role": "user", "content": "写脚本并运行"}],
            identity=UserIdentity("u1", "", "", "user"),
            raw_user_jwt="jwt",
            chat_id="c1",
            callback=callback,
        )
        assert answer == "验证完成"
        assert registry.exec_calls == 2
        assert [name for name, _ in registry.calls] == ["write_file", "exec", "write_file", "exec"]
        blocked_result = provider.requests[3]["messages"][-2]
        assert blocked_result["role"] == "tool"
        assert "unchanged retry" in blocked_result["content"]
    finally:
        await store.close()


async def test_duplicate_commands_in_one_model_response_execute_once(settings) -> None:
    batch = {
        "choices": [
            {
                "message": {
                    "role": "assistant",
                    "tool_calls": [
                        tool_call("1", "write_file", {"path": "sort.py", "content": "print('ok')"}),
                        tool_call("2", "exec", {"command": "python3 sort.py"}),
                        tool_call("3", "exec", {"command": "python3 sort.py"}),
                        tool_call("4", "exec", {"command": "python3 sort.py"}),
                    ],
                },
                "finish_reason": "tool_calls",
            }
        ]
    }
    final = {"choices": [{"message": {"role": "assistant", "content": "完成"}, "finish_reason": "stop"}]}
    provider = ScriptedProvider([batch, final])
    registry = RecordingRegistry()
    store = RuntimeStore(settings.database_url)
    await store.initialize()

    async def callback(event_type, payload):
        return None

    try:
        answer = await AgentLoop(provider, registry, store, max_tool_output_chars=10_000).run(
            spec=get_model_spec("work-light"),
            messages=[{"role": "user", "content": "写脚本并运行"}],
            identity=UserIdentity("u1", "", "", "user"),
            raw_user_jwt="jwt",
            chat_id="c1",
            callback=callback,
        )
        assert answer == "完成"
        assert [name for name, _ in registry.calls] == ["write_file", "exec"]
        duplicate_results = [
            message
            for message in provider.requests[1]["messages"]
            if message.get("role") == "tool" and "duplicate command" in message.get("content", "")
        ]
        assert len(duplicate_results) == 2
    finally:
        await store.close()


async def test_duplicate_identical_writes_execute_once(settings) -> None:
    batch = {
        "choices": [
            {
                "message": {
                    "role": "assistant",
                    "tool_calls": [
                        tool_call("1", "write_file", {"path": "sort.py", "content": "print('ok')"}),
                        tool_call("2", "write_file", {"path": "sort.py", "content": "print('ok')"}),
                        tool_call("3", "exec", {"command": "python3 sort.py"}),
                    ],
                },
                "finish_reason": "tool_calls",
            }
        ]
    }
    final = {"choices": [{"message": {"role": "assistant", "content": "完成"}, "finish_reason": "stop"}]}
    provider = ScriptedProvider([batch, final])
    registry = RecordingRegistry()
    store = RuntimeStore(settings.database_url)
    await store.initialize()

    async def callback(event_type, payload):
        return None

    try:
        answer = await AgentLoop(provider, registry, store, max_tool_output_chars=10_000).run(
            spec=get_model_spec("work-light"),
            messages=[{"role": "user", "content": "写脚本并运行"}],
            identity=UserIdentity("u1", "", "", "user"),
            raw_user_jwt="jwt",
            chat_id="c1",
            callback=callback,
        )
        assert answer == "完成"
        assert [name for name, _ in registry.calls] == ["write_file", "exec"]
        duplicate = provider.requests[1]["messages"][-2]
        assert duplicate["role"] == "tool"
        assert "same path and identical content" in duplicate["content"]
    finally:
        await store.close()


async def test_script_delivery_requires_successful_execution(settings) -> None:
    write = {
        "choices": [
            {
                "message": {
                    "role": "assistant",
                    "tool_calls": [tool_call("1", "write_file", {"path": "script.py", "content": "print('ok')"})],
                },
                "finish_reason": "tool_calls",
            }
        ]
    }
    read = {
        "choices": [
            {
                "message": {
                    "role": "assistant",
                    "tool_calls": [tool_call("2", "read_file", {"path": "script.py"})],
                },
                "finish_reason": "tool_calls",
            }
        ]
    }
    premature = {"choices": [{"message": {"role": "assistant", "content": "脚本完成"}, "finish_reason": "stop"}]}
    execute = {
        "choices": [
            {
                "message": {
                    "role": "assistant",
                    "tool_calls": [tool_call("3", "exec", {"command": "python3 script.py"})],
                },
                "finish_reason": "tool_calls",
            }
        ]
    }
    final = {"choices": [{"message": {"role": "assistant", "content": "脚本执行验证完成"}, "finish_reason": "stop"}]}
    provider = ScriptedProvider([write, read, premature, execute, final])
    registry = RecordingRegistry()
    store = RuntimeStore(settings.database_url)
    await store.initialize()

    async def callback(event_type, payload):
        return None

    try:
        answer = await AgentLoop(provider, registry, store, max_tool_output_chars=10_000).run(
            spec=get_model_spec("work-light"),
            messages=[{"role": "user", "content": "写一个 Python 脚本"}],
            identity=UserIdentity("u1", "", "", "user"),
            raw_user_jwt="jwt",
            chat_id="c1",
            callback=callback,
        )
        assert answer == "脚本执行验证完成"
        assert [name for name, _ in registry.calls] == ["write_file", "read_file", "exec"]
        assert "no command or process completed successfully" in provider.requests[3]["messages"][-1]["content"]
    finally:
        await store.close()


async def test_rejected_checkpoint_does_not_spin_without_tool_calls(settings) -> None:
    stopped = {"choices": [{"message": {"role": "assistant", "content": "已完成"}, "finish_reason": "stop"}]}
    provider = ScriptedProvider([stopped, stopped, stopped])
    store = RuntimeStore(settings.database_url)
    await store.initialize()
    events = []

    async def callback(event_type, payload):
        events.append(event_type)

    try:
        answer = await AgentLoop(
            provider,
            RecordingRegistry(),
            store,
            max_tool_output_chars=10_000,
        ).run(
            spec=get_model_spec("work-light"),
            messages=[{"role": "user", "content": "写一个报告"}],
            identity=UserIdentity("u1", "", "", "user"),
            raw_user_jwt="jwt",
            chat_id="c1",
            callback=callback,
        )
        assert answer.startswith("未完成：")
        assert len(provider.requests) == 3
        assert events.count("completion.rejected") == 2
        assert events.count("completion.unverified") == 1
    finally:
        await store.close()


async def test_bare_interactive_exec_is_blocked_before_registry(settings) -> None:
    first = {
        "choices": [
            {
                "message": {
                    "role": "assistant",
                    "tool_calls": [tool_call("1", "exec", {"command": "python3"})],
                },
                "finish_reason": "tool_calls",
            }
        ]
    }
    final = {"choices": [{"message": {"role": "assistant", "content": "无法执行"}, "finish_reason": "stop"}]}
    registry = RecordingRegistry()
    store = RuntimeStore(settings.database_url)
    await store.initialize()

    async def callback(event_type, payload):
        return None

    try:
        answer = await AgentLoop(
            ScriptedProvider([first, final]),
            registry,
            store,
            max_tool_output_chars=10_000,
        ).run(
            spec=get_model_spec("work-light"),
            messages=[{"role": "user", "content": "你好"}],
            identity=UserIdentity("u1", "", "", "user"),
            raw_user_jwt="jwt",
            chat_id="c1",
            callback=callback,
        )
        assert answer == "无法执行"
        assert registry.calls == []
    finally:
        await store.close()


def test_tool_events_render_one_friendly_completed_card() -> None:
    assert tool_event_details("tool.started", {"name": "exec"}) is None
    rendered = tool_event_details(
        "tool.completed",
        {
            "call_id": "call-1",
            "name": "exec",
            "arguments": {"command": "pytest -q"},
            "ok": True,
            "summary": "2 passed",
        },
    )
    assert rendered is not None
    assert 'name="执行命令"' in rendered
    assert "2 passed" in rendered
    assert 'done="false"' not in rendered