"""
Tests that when an LLM completion returns multiple tool calls including
sleep_until_next_trigger, the sleep call is ignored if other tools are present
so results can be processed in the next iteration.
"""
from unittest.mock import patch, MagicMock

from django.test import TestCase, tag
from django.contrib.auth import get_user_model

from api.models import (
    PersistentAgent,
    BrowserUseAgent,
    PersistentAgentCompletion,
    PersistentAgentKanbanCard,
    PersistentAgentToolCall,
    PersistentAgentStep,
    UserQuota,
)
from api.agent.tools.tool_manager import enable_tools


@tag("batch_event_parallel")
class TestBatchToolCallsWithSleep(TestCase):
    @classmethod
    def setUpTestData(cls):
        User = get_user_model()
        cls.user = User.objects.create_user(
            username='batchsleep@example.com', email='batchsleep@example.com', password='password'
        )
        quota, _ = UserQuota.objects.get_or_create(user=cls.user)
        quota.agent_limit = 100
        quota.save()

    def setUp(self):
        browser_agent = BrowserUseAgent.objects.create(user=self.user, name="browser-agent-for-batch-sleep-test")
        self.agent = PersistentAgent.objects.create(
            user=self.user,
            name="Batch Sleep Agent",
            charter="Test charter",
            browser_use_agent=browser_agent,
        )
        enable_tools(self.agent, ["sqlite_batch"])

    @patch('api.agent.core.event_processing._ensure_credit_for_tool', return_value={"cost": None, "credit": None})
    @patch('api.agent.core.event_processing.execute_enabled_tool', return_value={"status": "ok"})
    @patch('api.agent.core.event_processing.execute_update_charter', return_value={"status": "success"})
    @patch('api.agent.core.event_processing.execute_send_email', return_value={"status": "queued"})
    @patch('api.agent.core.event_processing.build_prompt_context')
    @patch('api.agent.core.event_processing._completion_with_failover')
    def test_batch_of_tools_ignores_sleep_when_others_present(
        self,
        mock_completion,
        mock_build_prompt,
        mock_send_email,
        mock_update_charter,
        mock_execute_enabled,
        _mock_credit,
    ):
        # Minimal prompt context and token usage
        mock_build_prompt.return_value = ([{"role": "system", "content": "sys"}, {"role": "user", "content": "go"}], 1000, None)

        # Construct four tool calls: send_email, update_charter, sqlite_batch, sleep
        def mk_tc(name, args):
            tc = MagicMock()
            tc.function = MagicMock()
            tc.function.name = name
            tc.function.arguments = args
            return tc

        tc_email = mk_tc('send_email', '{"to": "a@example.com", "subject": "hi", "html_body": "<p>Hi</p>"}')
        tc_charter = mk_tc('update_charter', '{"charter": "do x"}')
        tc_sqlite = mk_tc('sqlite_batch', '{"sql": "create table if not exists x(id int)"}')
        tc_sleep = mk_tc('sleep_until_next_trigger', '{}')

        msg = MagicMock()
        msg.tool_calls = [tc_email, tc_charter, tc_sqlite, tc_sleep]
        msg.content = None
        choice = MagicMock(); choice.message = msg
        resp = MagicMock(); resp.choices = [choice]
        # token usage dict present
        resp.model_extra = {"usage": MagicMock(prompt_tokens=10, completion_tokens=5, total_tokens=15, prompt_tokens_details=MagicMock(cached_tokens=0))}
        mock_completion.return_value = (resp, {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, "model": "m", "provider": "p"})

        # Run a single loop iteration
        from api.agent.core import event_processing as ep
        with patch.object(ep, 'MAX_AGENT_LOOP_ITERATIONS', 1):
            result_usage = ep._run_agent_loop(self.agent, is_first_run=False)

        # Validate DB records: 3 tool calls persisted and NO sleep step recorded
        calls = list(PersistentAgentToolCall.objects.all().order_by('step__created_at'))
        self.assertEqual(len(calls), 3)
        self.assertEqual([c.tool_name for c in calls], ['send_email', 'update_charter', 'sqlite_batch'])

        # Ensure no sleep step exists because sleep was ignored in mixed batch
        sleep_steps = PersistentAgentStep.objects.filter(description__icontains='sleep until next trigger')
        self.assertFalse(sleep_steps.exists(), "Sleep step should be ignored when other tools are present")

        # Assert token usage aggregated
        self.assertIn('total_tokens', result_usage)
        self.assertGreaterEqual(result_usage['total_tokens'], 15)

    @patch('api.agent.core.event_processing._ensure_credit_for_tool', return_value={"cost": None, "credit": None})
    @patch('api.agent.core.event_processing.execute_enabled_tool', return_value={"status": "ok", "auto_sleep_ok": True})
    @patch('api.agent.core.event_processing.build_prompt_context')
    @patch('api.agent.core.event_processing._completion_with_failover')
    def test_tool_call_dict_is_executed(
        self,
        mock_completion,
        mock_build_prompt,
        _mock_execute_enabled,
        _mock_credit,
    ):
        mock_build_prompt.return_value = (
            [{"role": "system", "content": "sys"}, {"role": "user", "content": "go"}],
            1000,
            None,
        )

        msg = MagicMock()
        msg.tool_calls = {
            "id": "call-1",
            "function": {"name": "sqlite_batch", "arguments": "{\"sql\": \"select 1\"}"},
        }
        msg.content = None

        choice = MagicMock()
        choice.message = msg
        resp = MagicMock()
        resp.choices = [choice]
        resp.model_extra = {
            "usage": MagicMock(
                prompt_tokens=10,
                completion_tokens=5,
                total_tokens=15,
                prompt_tokens_details=MagicMock(cached_tokens=0),
            )
        }

        mock_completion.return_value = (
            resp,
            {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, "model": "m", "provider": "p"},
        )

        from api.agent.core import event_processing as ep
        with patch.object(ep, 'MAX_AGENT_LOOP_ITERATIONS', 1):
            ep._run_agent_loop(self.agent, is_first_run=False)

        calls = list(PersistentAgentToolCall.objects.all().order_by('step__created_at'))
        self.assertEqual(len(calls), 1)
        self.assertEqual(calls[0].tool_name, 'sqlite_batch')

    @patch('api.agent.core.event_processing._ensure_credit_for_tool', return_value={"cost": None, "credit": None})
    @patch('api.agent.core.event_processing.execute_enabled_tool', return_value={"status": "ignored"})
    @patch('api.agent.core.event_processing.build_prompt_context')
    @patch('api.agent.core.event_processing._completion_with_failover')
    def test_sleep_only_batch_records_token_usage_once(
        self,
        mock_completion,
        mock_build_prompt,
        _mock_execute_enabled,
        _mock_credit,
    ):
        mock_build_prompt.return_value = ([{"role": "system", "content": "sys"}, {"role": "user", "content": "wait"}], 1000, None)

        def mk_sleep_call():
            tc = MagicMock()
            tc.function = MagicMock()
            tc.function.name = 'sleep_until_next_trigger'
            tc.function.arguments = '{}'
            return tc

        tc_sleep = mk_sleep_call()
        tc_sleep_followup = mk_sleep_call()

        msg = MagicMock()
        msg.tool_calls = [tc_sleep, tc_sleep_followup]
        msg.content = None
        choice = MagicMock(); choice.message = msg
        resp = MagicMock(); resp.choices = [choice]
        resp.model_extra = {"usage": MagicMock(prompt_tokens=10, completion_tokens=5, total_tokens=15, prompt_tokens_details=MagicMock(cached_tokens=0))}
        mock_completion.return_value = (resp, {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, "model": "m", "provider": "p"})

        from api.agent.core import event_processing as ep
        with patch.object(ep, 'MAX_AGENT_LOOP_ITERATIONS', 1):
            ep._run_agent_loop(self.agent, is_first_run=False)

        completions = list(PersistentAgentCompletion.objects.filter(agent=self.agent))
        self.assertEqual(len(completions), 1)
        completion = completions[0]
        self.assertEqual(completion.total_tokens, 15)
        sleep_steps = list(
            PersistentAgentStep.objects.filter(description__icontains='sleep until next trigger')
            .order_by('created_at')
        )
        self.assertEqual(len(sleep_steps), 2)
        for step in sleep_steps:
            self.assertEqual(step.completion_id, completion.id)

    @patch('api.agent.core.event_processing._ensure_credit_for_tool', return_value={"cost": None, "credit": None})
    @patch('api.agent.core.event_processing.execute_update_charter', return_value={"status": "ok", "auto_sleep_ok": True})
    @patch('api.agent.core.event_processing.execute_send_email', return_value={"status": "ok", "auto_sleep_ok": True})
    @patch('api.agent.core.event_processing.build_prompt_context')
    @patch('api.agent.core.event_processing._completion_with_failover')
    def test_successful_actions_short_circuit_to_sleep(
        self,
        mock_completion,
        mock_build_prompt,
        mock_send_email,
        mock_update_charter,
        _mock_credit,
    ):
        """A tool batch that opts-in to auto-sleep should end the loop immediately."""

        mock_build_prompt.return_value = ([{"role": "system", "content": "sys"}, {"role": "user", "content": "go"}], 1000, None)

        def mk_tc(name, args):
            tc = MagicMock()
            tc.function = MagicMock()
            tc.function.name = name
            tc.function.arguments = args
            return tc

        tc_email = mk_tc('send_email', '{"to": "a@example.com", "subject": "hi", "mobile_first_html": "<p>Hi</p>"}')
        tc_charter = mk_tc('update_charter', '{"new_charter": "Stay focused"}')
        tc_sleep = mk_tc('sleep_until_next_trigger', '{}')

        msg = MagicMock()
        msg.tool_calls = [tc_email, tc_charter, tc_sleep]
        msg.content = None
        choice = MagicMock(); choice.message = msg
        resp = MagicMock(); resp.choices = [choice]
        resp.model_extra = {"usage": MagicMock(prompt_tokens=10, completion_tokens=5, total_tokens=15, prompt_tokens_details=MagicMock(cached_tokens=0))}

        mock_completion.return_value = (resp, {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, "model": "m", "provider": "p"})

        from api.agent.core import event_processing as ep
        with patch.object(ep, 'MAX_AGENT_LOOP_ITERATIONS', 2):
            result_usage = ep._run_agent_loop(self.agent, is_first_run=False)

        # Only the initial completion should occur because the loop auto-sleeps
        self.assertEqual(mock_completion.call_count, 1)

        calls = list(PersistentAgentToolCall.objects.all().order_by('step__created_at'))
        self.assertEqual(len(calls), 2)
        self.assertEqual([c.tool_name for c in calls], ['send_email', 'update_charter'])

        # No explicit sleep step should exist because the loop short-circuited
        sleep_steps = PersistentAgentStep.objects.filter(description__icontains='sleep until next trigger')
        self.assertFalse(sleep_steps.exists())

        self.assertIn('total_tokens', result_usage)
        self.assertGreaterEqual(result_usage['total_tokens'], 15)

    @patch('api.agent.core.event_processing._ensure_credit_for_tool', return_value={"cost": None, "credit": None})
    @patch('api.agent.core.event_processing.execute_spawn_web_task', return_value={"status": "pending", "auto_sleep_ok": True})
    @patch('api.agent.core.event_processing.execute_send_email', return_value={"status": "sent", "auto_sleep_ok": True})
    @patch('api.agent.core.event_processing.build_prompt_context')
    @patch('api.agent.core.event_processing._completion_with_failover')
    def test_auto_sleep_waits_for_all_tool_calls(
        self,
        mock_completion,
        mock_build_prompt,
        mock_send_email,
        mock_spawn_task,
        _mock_credit,
    ):
        """Ensure we execute every actionable tool call before honoring auto-sleep."""

        mock_build_prompt.return_value = ([{"role": "system", "content": "sys"}, {"role": "user", "content": "go"}], 1000, None)

        def mk_tc(name, args):
            tc = MagicMock()
            tc.function = MagicMock()
            tc.function.name = name
            tc.function.arguments = args
            return tc

        tc_email = mk_tc('send_email', '{"to": "a@example.com", "subject": "hi", "mobile_first_html": "<p>Hi</p>"}')
        tc_spawn = mk_tc('spawn_web_task', '{"url": "https://example.com", "charter": "do something"}')

        msg = MagicMock()
        msg.tool_calls = [tc_email, tc_spawn]
        msg.content = None
        choice = MagicMock(); choice.message = msg
        resp = MagicMock(); resp.choices = [choice]
        resp.model_extra = {"usage": MagicMock(prompt_tokens=10, completion_tokens=5, total_tokens=15, prompt_tokens_details=MagicMock(cached_tokens=0))}

        msg_followup = MagicMock()
        msg_followup.tool_calls = None
        msg_followup.function_call = None
        msg_followup.content = "Done"
        followup_choice = MagicMock(); followup_choice.message = msg_followup
        followup_resp = MagicMock(); followup_resp.choices = [followup_choice]
        followup_resp.model_extra = {"usage": MagicMock(prompt_tokens=4, completion_tokens=2, total_tokens=6, prompt_tokens_details=MagicMock(cached_tokens=0))}

        mock_completion.side_effect = [
            (resp, {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, "model": "m", "provider": "p"}),
            (followup_resp, {"prompt_tokens": 4, "completion_tokens": 2, "total_tokens": 6, "model": "m", "provider": "p"}),
        ]

        from api.agent.core import event_processing as ep
        with patch.object(ep, 'MAX_AGENT_LOOP_ITERATIONS', 2):
            ep._run_agent_loop(self.agent, is_first_run=False)

        # Auto-sleep should fire after both actionable calls run, so only one completion is needed.
        self.assertEqual(mock_completion.call_count, 1)

        mock_send_email.assert_called_once()
        mock_spawn_task.assert_called_once()

        calls = list(PersistentAgentToolCall.objects.all().order_by('step__created_at'))
        self.assertEqual(len(calls), 2)
        self.assertEqual([c.tool_name for c in calls], ['send_email', 'spawn_web_task'])

        sleep_steps = PersistentAgentStep.objects.filter(description__icontains='sleep until next trigger')
        self.assertFalse(sleep_steps.exists())

    @patch('api.agent.core.event_processing._ensure_credit_for_tool', return_value={"cost": None, "credit": None})
    @patch('api.agent.core.event_processing.execute_update_charter', return_value={"status": "ok", "auto_sleep_ok": True})
    @patch('api.agent.core.event_processing.execute_send_email', return_value={"status": "ok", "auto_sleep_ok": True})
    @patch('api.agent.core.event_processing.build_prompt_context')
    @patch('api.agent.core.event_processing._completion_with_failover')
    def test_auto_sleep_triggers_without_sleep_tool_call(self, mock_completion, mock_build_prompt, *_mocks):
        """Auto-sleep should trigger when all tool results allow it, even with no explicit sleep call."""

        mock_build_prompt.return_value = ([{"role": "system", "content": "sys"}, {"role": "user", "content": "go"}], 1000, None)

        def mk_tc(name, args):
            tc = MagicMock()
            tc.function = MagicMock()
            tc.function.name = name
            tc.function.arguments = args
            return tc

        tc_email = mk_tc('send_email', '{"to": "a@example.com", "subject": "hi", "mobile_first_html": "<p>Hi</p>"}')
        tc_charter = mk_tc('update_charter', '{"new_charter": "Stay focused"}')

        msg = MagicMock()
        msg.tool_calls = [tc_email, tc_charter]
        msg.content = None

        choice = MagicMock(); choice.message = msg
        resp = MagicMock(); resp.choices = [choice]
        resp.model_extra = {"usage": MagicMock(prompt_tokens=10, completion_tokens=5, total_tokens=15, prompt_tokens_details=MagicMock(cached_tokens=0))}

        msg_followup = MagicMock()
        msg_followup.tool_calls = None
        msg_followup.function_call = None
        msg_followup.content = "Done"
        followup_choice = MagicMock(); followup_choice.message = msg_followup
        followup_resp = MagicMock(); followup_resp.choices = [followup_choice]
        followup_resp.model_extra = {"usage": MagicMock(prompt_tokens=4, completion_tokens=2, total_tokens=6, prompt_tokens_details=MagicMock(cached_tokens=0))}

        mock_completion.side_effect = [
            (resp, {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, "model": "m", "provider": "p"}),
            (followup_resp, {"prompt_tokens": 4, "completion_tokens": 2, "total_tokens": 6, "model": "m", "provider": "p"}),
        ]

        from api.agent.core import event_processing as ep
        with patch.object(ep, 'MAX_AGENT_LOOP_ITERATIONS', 2):
            ep._run_agent_loop(self.agent, is_first_run=False)

        # The loop should stop after the first completion because both tools allowed auto-sleep.
        self.assertEqual(mock_completion.call_count, 1)

        # The actionable tool calls were still executed and recorded once each
        calls = list(PersistentAgentToolCall.objects.all().order_by('step__created_at'))
        self.assertEqual(len(calls), 2)
        self.assertEqual([c.tool_name for c in calls], ['send_email', 'update_charter'])

        sleep_steps = PersistentAgentStep.objects.filter(description__icontains='sleep until next trigger')
        self.assertFalse(sleep_steps.exists())

    @patch('api.agent.core.event_processing._ensure_credit_for_tool', return_value={"cost": None, "credit": None})
    @patch('api.agent.core.event_processing.execute_update_charter', return_value={"status": "ok", "auto_sleep_ok": True})
    @patch('api.agent.core.event_processing.execute_send_email', return_value={"status": "ok", "auto_sleep_ok": True})
    @patch('api.agent.core.event_processing.build_prompt_context')
    @patch('api.agent.core.event_processing._completion_with_failover')
    def test_auto_sleep_not_overridden_with_open_plan_work(
        self,
        mock_completion,
        mock_build_prompt,
        *_mocks,
    ):
        notices = []

        def build_prompt_side_effect(*_args, **kwargs):
            notices.append(kwargs.get("continuation_notice"))
            return ([{"role": "system", "content": "sys"}, {"role": "user", "content": "go"}], 1000, None)

        mock_build_prompt.side_effect = build_prompt_side_effect

        self.agent.schedule = "0 9 * * *"
        self.agent.save(update_fields=["schedule"])

        PersistentAgentKanbanCard.objects.create(
            assigned_agent=self.agent,
            title="Finish outstanding work",
            status=PersistentAgentKanbanCard.Status.TODO,
        )

        def mk_tc(name, args):
            tc = MagicMock()
            tc.function = MagicMock()
            tc.function.name = name
            tc.function.arguments = args
            return tc

        tc_email = mk_tc('send_email', '{"to": "a@example.com", "subject": "hi", "mobile_first_html": "<p>Hi</p>"}')
        tc_charter = mk_tc('update_charter', '{"new_charter": "Stay focused"}')

        msg = MagicMock()
        msg.tool_calls = [tc_email, tc_charter]
        msg.content = None

        choice = MagicMock(); choice.message = msg
        resp = MagicMock(); resp.choices = [choice]
        resp.model_extra = {"usage": MagicMock(prompt_tokens=10, completion_tokens=5, total_tokens=15, prompt_tokens_details=MagicMock(cached_tokens=0))}

        mock_completion.side_effect = [
            (resp, {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, "model": "m", "provider": "p"}),
        ]

        from api.agent.core import event_processing as ep
        with patch.object(ep, 'MAX_AGENT_LOOP_ITERATIONS', 2):
            ep._run_agent_loop(self.agent, is_first_run=False)

        self.assertEqual(mock_completion.call_count, 1)
        self.assertEqual(len(notices), 1)
        self.assertIsNone(notices[0])

    @patch('api.agent.core.event_processing._ensure_credit_for_tool', return_value={"cost": None, "credit": None})
    @patch('api.agent.core.event_processing.execute_send_chat_message', return_value={"status": "ok", "auto_sleep_ok": True})
    @patch('api.agent.core.event_processing.build_prompt_context')
    @patch('api.agent.core.event_processing._completion_with_failover')
    def test_message_with_auto_sleep_ok_stops(
        self,
        mock_completion,
        mock_build_prompt,
        mock_send_chat,
        _mock_credit,
    ):
        """When a message is sent with auto_sleep_ok=True (will_continue_work=false), agent stops."""
        mock_build_prompt.return_value = (
            [{"role": "system", "content": "sys"}, {"role": "user", "content": "go"}],
            1000,
            None,
        )

        def mk_tc(name, args):
            tc = MagicMock()
            tc.function = MagicMock()
            tc.function.name = name
            tc.function.arguments = args
            return tc

        # will_continue_work=false means auto_sleep_ok=True from send_chat_message
        tc_message = mk_tc('send_chat_message', '{"body": "Done.", "will_continue_work": false}')

        msg = MagicMock()
        msg.tool_calls = [tc_message]
        msg.content = None
        choice = MagicMock(); choice.message = msg
        resp = MagicMock(); resp.choices = [choice]
        resp.model_extra = {"usage": MagicMock(prompt_tokens=10, completion_tokens=5, total_tokens=15, prompt_tokens_details=MagicMock(cached_tokens=0))}

        mock_completion.side_effect = [
            (resp, {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, "model": "m", "provider": "p"}),
        ]

        from api.agent.core import event_processing as ep
        with patch.object(ep, 'MAX_AGENT_LOOP_ITERATIONS', 2):
            ep._run_agent_loop(self.agent, is_first_run=False)

        self.assertEqual(mock_completion.call_count, 1)
        mock_send_chat.assert_called_once()

    @patch('api.agent.core.event_processing._ensure_credit_for_tool', return_value={"cost": None, "credit": None})
    @patch('api.agent.core.event_processing.execute_update_plan', return_value={"status": "ok"})
    @patch('api.agent.core.event_processing.execute_enabled_tool', return_value={"status": "ok", "auto_sleep_ok": True})
    @patch('api.agent.core.event_processing.execute_send_chat_message', return_value={"status": "ok", "auto_sleep_ok": False})
    @patch('api.agent.core.event_processing.build_prompt_context')
    @patch('api.agent.core.event_processing._completion_with_failover')
    def test_explicit_stop_overrides_prior_continue(
        self,
        mock_completion,
        mock_build_prompt,
        mock_send_chat,
        mock_execute_enabled,
        mock_update_plan,
        _mock_credit,
    ):
        """Stop when the final tool call explicitly sets will_continue_work=false."""
        mock_build_prompt.return_value = (
            [{"role": "system", "content": "sys"}, {"role": "user", "content": "go"}],
            1000,
            None,
        )

        def mk_tc(name, args):
            tc = MagicMock()
            tc.function = MagicMock()
            tc.function.name = name
            tc.function.arguments = args
            return tc

        tc_message = mk_tc('send_chat_message', '{"body": "Report done.", "will_continue_work": true}')
        tc_plan = mk_tc(
            'update_plan',
            '{"plan": [{"step": "Deliver final report", "status": "done"}], "will_continue_work": false}'
        )

        msg = MagicMock()
        msg.tool_calls = [tc_message, tc_plan]
        msg.content = None
        choice = MagicMock(); choice.message = msg
        resp = MagicMock(); resp.choices = [choice]
        resp.model_extra = {"usage": MagicMock(prompt_tokens=10, completion_tokens=5, total_tokens=15, prompt_tokens_details=MagicMock(cached_tokens=0))}

        msg_followup = MagicMock()
        msg_followup.tool_calls = None
        msg_followup.function_call = None
        msg_followup.content = "Extra turn"
        followup_choice = MagicMock(); followup_choice.message = msg_followup
        followup_resp = MagicMock(); followup_resp.choices = [followup_choice]
        followup_resp.model_extra = {"usage": MagicMock(prompt_tokens=4, completion_tokens=2, total_tokens=6, prompt_tokens_details=MagicMock(cached_tokens=0))}

        mock_completion.side_effect = [
            (resp, {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, "model": "m", "provider": "p"}),
            (followup_resp, {"prompt_tokens": 4, "completion_tokens": 2, "total_tokens": 6, "model": "m", "provider": "p"}),
        ]

        from api.agent.core import event_processing as ep
        with patch.object(ep, 'MAX_AGENT_LOOP_ITERATIONS', 2):
            ep._run_agent_loop(self.agent, is_first_run=False)

        self.assertEqual(mock_completion.call_count, 1)
        mock_send_chat.assert_called_once()
        mock_update_plan.assert_called_once()
        mock_execute_enabled.assert_not_called()
