import asyncio
import json
import sys
import uuid
from unittest.mock import AsyncMock, MagicMock, mock_open, patch

import pytest

from control_backend.agents.bdi.bdi_program_manager import BDIProgramManager
from control_backend.core.agent_system import InternalMessage
from control_backend.schemas.program import (
    BasicNorm,
    ConditionalNorm,
    Goal,
    InferredBelief,
    KeywordBelief,
    Phase,
    Plan,
    Program,
    Trigger,
)

# Fix Windows Proactor loop for zmq
if sys.platform.startswith("win"):
    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())


def make_valid_program_json(norm="N1", goal="G1") -> str:
    return Program(
        phases=[
            Phase(
                id=uuid.uuid4(),
                name="Basic Phase",
                norms=[
                    BasicNorm(
                        id=uuid.uuid4(),
                        name=norm,
                        norm=norm,
                    ),
                ],
                goals=[
                    Goal(
                        id=uuid.uuid4(),
                        name=goal,
                        description="This description can be used to determine whether the goal "
                        "has been achieved.",
                        plan=Plan(
                            id=uuid.uuid4(),
                            name="Goal Plan",
                            steps=[],
                        ),
                        can_fail=False,
                    ),
                ],
                triggers=[],
            ),
        ],
    ).model_dump_json()


@pytest.mark.asyncio
async def test_create_agentspeak_and_send_to_bdi(mock_settings):
    manager = BDIProgramManager(name="program_manager_test")
    manager.send = AsyncMock()

    program = Program.model_validate_json(make_valid_program_json())

    with patch("builtins.open", mock_open()) as mock_file:
        await manager._create_agentspeak_and_send_to_bdi(program)

        # Check file writing
        mock_file.assert_called_with(mock_settings.behaviour_settings.agentspeak_file, "w")
        handle = mock_file()
        handle.write.assert_called()

    assert manager.send.await_count == 1
    msg: InternalMessage = manager.send.await_args[0][0]
    assert msg.thread == "new_program"
    assert msg.to == mock_settings.agent_settings.bdi_core_name
    assert msg.body == mock_settings.behaviour_settings.agentspeak_file


@pytest.mark.asyncio
async def test_receive_programs_valid_and_invalid():
    sub = AsyncMock()
    sub.recv_multipart.side_effect = [
        (b"program", b"{bad json"),
        (b"program", make_valid_program_json().encode()),
    ]

    manager = BDIProgramManager(name="program_manager_test")
    manager._internal_pub_socket = AsyncMock()
    manager.sub_socket = sub
    manager._create_agentspeak_and_send_to_bdi = AsyncMock()
    manager._send_clear_llm_history = AsyncMock()
    manager._send_program_to_user_interrupt = AsyncMock()
    manager._send_beliefs_to_semantic_belief_extractor = AsyncMock()
    manager._send_goals_to_semantic_belief_extractor = AsyncMock()

    try:
        # Will give StopAsyncIteration when the predefined `sub.recv_multipart` side-effects run out
        await manager._receive_programs()
    except StopAsyncIteration:
        pass

    # Only valid Program should have triggered _send_to_bdi
    assert manager._create_agentspeak_and_send_to_bdi.await_count == 1
    forwarded: Program = manager._create_agentspeak_and_send_to_bdi.await_args[0][0]
    assert forwarded.phases[0].norms[0].name == "N1"
    assert forwarded.phases[0].goals[0].name == "G1"

    # Verify history clear was triggered exactly once (for the valid program)
    # The invalid program loop `continue`s before calling _send_clear_llm_history
    assert manager._send_clear_llm_history.await_count == 1


@pytest.mark.asyncio
async def test_send_clear_llm_history(mock_settings):
    # Ensure the mock returns a string for the agent name (just like in your LLM tests)
    mock_settings.agent_settings.llm_agent_name = "llm_agent"

    manager = BDIProgramManager(name="program_manager_test")
    manager.send = AsyncMock()

    await manager._send_clear_llm_history()

    assert manager.send.await_count == 2
    msg: InternalMessage = manager.send.await_args_list[0][0][0]

    # Verify the content and recipient
    assert msg.body == "clear_history"


@pytest.mark.asyncio
async def test_handle_message_transition_phase(mock_settings):
    mock_settings.agent_settings.user_interrupt_name = "user_interrupt_agent"
    manager = BDIProgramManager(name="program_manager_test")
    manager.send = AsyncMock()

    # Setup state
    prog = Program.model_validate_json(make_valid_program_json(norm="N1", goal="G1"))
    manager._initialize_internal_state(prog)

    # Test valid transition (to same phase for simplicity, or we need 2 phases)
    # Let's create a program with 2 phases
    phase2_id = uuid.uuid4()
    phase2 = Phase(id=phase2_id, name="Phase 2", norms=[], goals=[], triggers=[])
    prog.phases.append(phase2)
    manager._initialize_internal_state(prog)

    current_phase_id = str(prog.phases[0].id)
    next_phase_id = str(phase2_id)

    payload = json.dumps({"old": current_phase_id, "new": next_phase_id})
    msg = InternalMessage(to="me", sender="bdi", body=payload, thread="transition_phase")

    await manager.handle_message(msg)

    assert str(manager._phase.id) == next_phase_id

    # Allow background tasks to run (add_behavior)
    await asyncio.sleep(0)

    # Check notifications sent
    # 1. beliefs to extractor
    # 2. goals to extractor
    # 3. notification to user interrupt

    assert manager.send.await_count >= 3

    # Verify user interrupt notification
    calls = manager.send.await_args_list
    ui_msgs = [
        c[0][0] for c in calls if c[0][0].to == mock_settings.agent_settings.user_interrupt_name
    ]
    assert len(ui_msgs) > 0
    assert ui_msgs[-1].body == next_phase_id


@pytest.mark.asyncio
async def test_handle_message_transition_phase_desync():
    manager = BDIProgramManager(name="program_manager_test")
    manager.logger = MagicMock()

    prog = Program.model_validate_json(make_valid_program_json())
    manager._initialize_internal_state(prog)

    current_phase_id = str(prog.phases[0].id)

    # Request transition from WRONG old phase
    payload = json.dumps({"old": "wrong_id", "new": "some_new_id"})
    msg = InternalMessage(to="me", sender="bdi", body=payload, thread="transition_phase")

    await manager.handle_message(msg)

    # Should warn and do nothing
    manager.logger.warning.assert_called_once()
    assert "Phase transition desync detected" in manager.logger.warning.call_args[0][0]
    assert str(manager._phase.id) == current_phase_id


@pytest.mark.asyncio
async def test_handle_message_transition_phase_end(mock_settings):
    mock_settings.agent_settings.user_interrupt_name = "user_interrupt_agent"
    manager = BDIProgramManager(name="program_manager_test")
    manager.send = AsyncMock()

    prog = Program.model_validate_json(make_valid_program_json())
    manager._initialize_internal_state(prog)
    current_phase_id = str(prog.phases[0].id)

    payload = json.dumps({"old": current_phase_id, "new": "end"})
    msg = InternalMessage(to="me", sender="bdi", body=payload, thread="transition_phase")

    await manager.handle_message(msg)

    assert manager._phase is None

    # Allow background tasks to run (add_behavior)
    await asyncio.sleep(0)

    # Verify notification to user interrupt
    assert manager.send.await_count == 1
    msg_sent = manager.send.await_args[0][0]
    assert msg_sent.to == mock_settings.agent_settings.user_interrupt_name
    assert msg_sent.body == "end"


@pytest.mark.asyncio
async def test_handle_message_achieve_goal(mock_settings):
    mock_settings.agent_settings.text_belief_extractor_name = "text_belief_extractor_agent"
    manager = BDIProgramManager(name="program_manager_test")
    manager.send = AsyncMock()

    prog = Program.model_validate_json(make_valid_program_json(goal="TargetGoal"))
    manager._initialize_internal_state(prog)

    goal_id = str(prog.phases[0].goals[0].id)

    msg = InternalMessage(to="me", sender="ui", body=goal_id, thread="achieve_goal")

    await manager.handle_message(msg)

    # Should send achieved goals to text extractor
    assert manager.send.await_count == 1
    msg_sent = manager.send.await_args[0][0]
    assert msg_sent.to == mock_settings.agent_settings.text_belief_extractor_name
    assert msg_sent.thread == "achieved_goals"

    # Verify body
    from control_backend.schemas.belief_list import GoalList

    gl = GoalList.model_validate_json(msg_sent.body)
    assert len(gl.goals) == 1
    assert gl.goals[0].name == "TargetGoal"


@pytest.mark.asyncio
async def test_handle_message_achieve_goal_not_found():
    manager = BDIProgramManager(name="program_manager_test")
    manager.send = AsyncMock()
    manager.logger = MagicMock()

    prog = Program.model_validate_json(make_valid_program_json())
    manager._initialize_internal_state(prog)

    msg = InternalMessage(to="me", sender="ui", body="non_existent_id", thread="achieve_goal")

    await manager.handle_message(msg)

    manager.send.assert_not_called()
    manager.logger.debug.assert_called()


@pytest.mark.asyncio
async def test_setup(mock_settings):
    manager = BDIProgramManager(name="program_manager_test")
    manager.send = AsyncMock()

    def close_coro(coro):
        coro.close()
        return MagicMock()

    manager.add_behavior = MagicMock(side_effect=close_coro)

    mock_context = MagicMock()
    mock_sub = MagicMock()
    mock_context.socket.return_value = mock_sub

    with patch(
        "control_backend.agents.bdi.bdi_program_manager.Context.instance", return_value=mock_context
    ):
        # We also need to mock file writing in _create_agentspeak_and_send_to_bdi
        with patch("builtins.open", new_callable=MagicMock):
            await manager.setup()

    # Check logic
    # 1. Sends default empty program to BDI
    assert manager.send.await_count == 1
    assert manager.send.await_args[0][0].to == mock_settings.agent_settings.bdi_core_name

    # 2. Connects SUB socket
    mock_sub.connect.assert_called_with(mock_settings.zmq_settings.internal_sub_address)
    mock_sub.subscribe.assert_called_with("program")

    # 3. Adds behavior
    manager.add_behavior.assert_called()


@pytest.mark.asyncio
async def test_send_program_to_user_interrupt(mock_settings):
    """Test directly sending the program to the user interrupt agent."""
    mock_settings.agent_settings.user_interrupt_name = "user_interrupt_agent"

    manager = BDIProgramManager(name="program_manager_test")
    manager.send = AsyncMock()

    program = Program.model_validate_json(make_valid_program_json())

    await manager._send_program_to_user_interrupt(program)

    assert manager.send.await_count == 1
    msg = manager.send.await_args[0][0]
    assert msg.to == "user_interrupt_agent"
    assert msg.thread == "new_program"
    assert "Basic Phase" in msg.body


@pytest.mark.asyncio
async def test_complex_program_extraction():
    manager = BDIProgramManager(name="program_manager_test")

    # 1. Create Complex Components

    # Inferred Belief (A & B)
    belief_left = KeywordBelief(id=uuid.uuid4(), name="b1", keyword="hot")
    belief_right = KeywordBelief(id=uuid.uuid4(), name="b2", keyword="sunny")
    inferred_belief = InferredBelief(
        id=uuid.uuid4(), name="b_inf", operator="AND", left=belief_left, right=belief_right
    )

    # Conditional Norm
    cond_norm = ConditionalNorm(
        id=uuid.uuid4(), name="norm_cond", norm="wear_hat", condition=inferred_belief
    )

    # Trigger with Inferred Belief condition
    dummy_plan = Plan(id=uuid.uuid4(), name="dummy_plan", steps=[])
    trigger = Trigger(id=uuid.uuid4(), name="trigger_1", condition=inferred_belief, plan=dummy_plan)

    #  Nested Goal
    sub_goal = Goal(
        id=uuid.uuid4(),
        name="sub_goal",
        description="desc",
        plan=Plan(id=uuid.uuid4(), name="empty", steps=[]),
        can_fail=True,
    )

    parent_goal = Goal(
        id=uuid.uuid4(),
        name="parent_goal",
        description="desc",
        # The plan contains the sub_goal as a step
        plan=Plan(id=uuid.uuid4(), name="parent_plan", steps=[sub_goal]),
        can_fail=False,
    )

    # 2. Assemble Program
    phase = Phase(
        id=uuid.uuid4(),
        name="Complex Phase",
        norms=[cond_norm],
        goals=[parent_goal],
        triggers=[trigger],
    )
    program = Program(phases=[phase])

    # 3. Initialize Internal State (Triggers _populate_goal_mapping -> Nested Goal logic)
    manager._initialize_internal_state(program)

    # Assertion for Line 53-54 (Mapping population)
    # Both parent and sub-goal should be mapped
    assert str(parent_goal.id) in manager._goal_mapping
    assert str(sub_goal.id) in manager._goal_mapping

    # 4. Test Belief Extraction (Triggers lines 132-133, 142-146)
    beliefs = manager._extract_current_beliefs()

    # Should extract recursive beliefs from cond_norm and trigger
    # Inferred belief splits into Left + Right. Since we use it twice, we get duplicates
    # checking existence is enough.
    belief_names = [b.name for b in beliefs]
    assert "b1" in belief_names
    assert "b2" in belief_names

    # 5. Test Goal Extraction (Triggers lines 173, 185)
    goals = manager._extract_current_goals()

    goal_names = [g.name for g in goals]
    assert "parent_goal" in goal_names
    assert "sub_goal" in goal_names