from __future__ import annotations

import textwrap
from typing import Any


def _sanitize_id(raw: str) -> str:
    """Turn an agent id like 'my-agent' into a valid Python identifier suffix."""
    return raw.replace("-", "_").replace(".", "_").replace("/", "_")


def _escape_str(raw: str) -> str:
    """Escape a string for safe interpolation into double-quoted Python string literals."""
    return raw.replace("\\", "\\\\").replace('"', '\\"')


def _build_single_agent_test(agent_id: str) -> str:
    safe = _sanitize_id(agent_id)
    esc = _escape_str(agent_id)
    return textwrap.dedent(f"""\

        def test_{safe}_emits_events() -> None:
            mock = MockAmplitudeAI()
            agent = mock.agent("{esc}")
            with agent.session(user_id="verify-user", session_id="verify-session") as s:
                s.track_user_message("test message")
            events = mock.get_events()
            assert len(events) >= 2
            mock.assert_event_tracked("[Agent] User Message", agent_id="{esc}")
            mock.assert_session_closed("verify-session")
    """)


def _build_multi_agent_test(parent_id: str, child_id: str) -> str:
    esc_p, esc_c = _escape_str(parent_id), _escape_str(child_id)
    return textwrap.dedent(f"""\

        def test_multi_agent_delegation() -> None:
            mock = MockAmplitudeAI()
            parent = mock.agent("{esc_p}")
            child = parent.child("{esc_c}")
            with parent.session(user_id="verify-user", session_id="verify-multi") as s:
                with s.run_as(child) as cs:
                    cs.track_user_message("delegated task")
            child_events = mock.events_for_agent("{esc_c}")
            assert len(child_events) > 0
            assert child_events[0]["event_properties"]["[Agent] Session ID"] == "verify-multi"
    """)


def _build_async_multi_agent_test(parent_id: str, child_id: str) -> str:
    esc_p, esc_c = _escape_str(parent_id), _escape_str(child_id)
    return textwrap.dedent(f"""\

        @pytest.mark.asyncio
        async def test_async_multi_agent() -> None:
            mock = MockAmplitudeAI()
            parent = mock.agent("{esc_p}")
            child = parent.child("{esc_c}")
            async with parent.session(user_id="verify-user", session_id="verify-async") as s:
                async with s.arun_as(child) as cs:
                    cs.track_user_message("async delegated task")
            child_events = mock.events_for_agent("{esc_c}")
            assert len(child_events) > 0
            assert child_events[0]["event_properties"]["[Agent] Session ID"] == "verify-async"
    """)


def _build_tool_attribution_test(parent_id: str, child_id: str) -> str:
    safe_child = _sanitize_id(child_id)
    esc_p, esc_c = _escape_str(parent_id), _escape_str(child_id)
    return textwrap.dedent(f"""\

        def test_tool_calls_attributed_to_child_{safe_child}() -> None:
            mock = MockAmplitudeAI()
            parent = mock.agent("{esc_p}")
            child = parent.child("{esc_c}")
            with parent.session(user_id="verify-user", session_id="verify-tools") as s:
                with s.run_as(child) as cs:
                    cs.track_tool_call("search_knowledge_base", 150, True)
            tool_events = [e for e in mock.get_events() if e["event_type"] == "[Agent] Tool Call"]
            assert len(tool_events) == 1
            assert tool_events[0]["event_properties"]["[Agent] Agent ID"] == "{esc_c}"
            assert tool_events[0]["event_properties"]["[Agent] Parent Agent ID"] == "{esc_p}"
    """)


def _build_error_recovery_test(parent_id: str, child_id: str) -> str:
    esc_p, esc_c = _escape_str(parent_id), _escape_str(child_id)
    return textwrap.dedent(f"""\

        def test_run_as_restores_parent_context_after_error() -> None:
            mock = MockAmplitudeAI()
            parent = mock.agent("{esc_p}")
            faulty_child = parent.child("{esc_c}-faulty")
            with parent.session(user_id="verify-user", session_id="verify-error") as s:
                try:
                    with s.run_as(faulty_child) as cs:
                        raise RuntimeError("child failed")
                except RuntimeError:
                    s.track_user_message("recovering from child failure")
            parent_events = mock.events_for_agent("{esc_p}")
            recovery_msgs = [e for e in parent_events if e["event_type"] == "[Agent] User Message"]
            assert len(recovery_msgs) > 0
            assert recovery_msgs[0]["event_properties"]["[Agent] Agent ID"] == "{esc_p}"
    """)


def generate_verify_test(scan_result: dict[str, Any]) -> str:
    """Generate a complete pytest verification test file from a scan_project result."""
    agents: list[dict[str, Any]] = scan_result.get("agents", [])
    is_multi_agent: bool = scan_result.get("is_multi_agent", False)
    multi_agent_signals: list[str] = scan_result.get("multi_agent_signals", [])

    should_emit_multi_agent_tests = (
        is_multi_agent or len(multi_agent_signals) > 0
    ) and len(agents) >= 2

    imports = [
        "import pytest",
        "",
        "from amplitude_ai.testing import MockAmplitudeAI",
    ]

    lines: list[str] = imports[:]

    for agent in agents:
        agent_id = agent.get("inferred_id", "default")
        lines.append(_build_single_agent_test(agent_id))

    if should_emit_multi_agent_tests:
        parent_id = agents[0].get("inferred_id", "parent")
        child_id = agents[1].get("inferred_id", "child")
        lines.append(_build_multi_agent_test(parent_id, child_id))
        lines.append(_build_async_multi_agent_test(parent_id, child_id))
        lines.append(_build_tool_attribution_test(parent_id, child_id))
        lines.append(_build_error_recovery_test(parent_id, child_id))
    elif (is_multi_agent or len(multi_agent_signals) > 0) and agents:
        parent_id = agents[0].get("inferred_id", "parent")
        child_id = f"{parent_id}-child"
        lines.append(_build_multi_agent_test(parent_id, child_id))
        lines.append(_build_async_multi_agent_test(parent_id, child_id))
        lines.append(_build_tool_attribution_test(parent_id, child_id))
        lines.append(_build_error_recovery_test(parent_id, child_id))

    return "\n".join(lines).rstrip() + "\n"
