"""Tests for spawn_web_task child-count tracking.

Validates that parallel spawned tasks do not affect recursion depth
argument but correctly increment the outstanding-children counter in
Redis so the agent loop keeps the cycle open while children run.
"""

from datetime import timedelta
import threading
import time
from itertools import count
from unittest.mock import patch, MagicMock
from django.test import TransactionTestCase, override_settings, tag
from django.contrib.auth import get_user_model
from django.utils import timezone

from api.encryption import SecretsEncryption
from api.models import (
    PersistentAgent,
    BrowserUseAgent,
    BrowserUseAgentTask,
    BrowserConfig,
    EvalRun,
    PersistentAgentSecret,
    TaskCredit,
)
from api.agent.core.budget import AgentBudgetManager, BudgetContext, set_current_context as set_budget_context
from api.agent.tools.spawn_web_task import execute_spawn_web_task
from constants.grant_types import GrantTypeChoices
from constants.plans import PlanNames, PlanNamesChoices
from util.analytics import AnalyticsEvent


@tag("batch_spawn_depth")
class SpawnDepthTrackingTests(TransactionTestCase):
    """Test that parallel spawn_web_task calls correctly track depth."""
    
    def setUp(self):
        """Set up test data."""
        User = get_user_model()
        self.user = User.objects.create_user(
            username='test@example.com',
            email='test@example.com'
        )
        
        # Create a browser agent
        with patch.object(BrowserUseAgent, 'select_random_proxy', return_value=None):
            self.browser_agent = BrowserUseAgent.objects.create(
                user=self.user,
                name="Test Browser Agent"
            )
        
        # Create a persistent agent
        self.agent = PersistentAgent.objects.create(
            user=self.user,
            name="Test Agent",
            charter="Test",
            browser_use_agent=self.browser_agent
        )
        TaskCredit.objects.create(
            user=self.user,
            credits=50,
            granted_date=timezone.now(),
            expiration_date=timezone.now() + timedelta(days=30),
            plan=PlanNamesChoices.FREE,
            grant_type=GrantTypeChoices.PROMO,
        )
        
        # Initialize a budget cycle
        self.budget_id, _, _ = AgentBudgetManager.find_or_start_cycle(
            agent_id=str(self.agent.id),
            max_steps=100,
            max_depth=3  # Allow depth up to 3 for testing
        )
        
        # Create initial branch at depth 0
        self.branch_id = AgentBudgetManager.create_branch(
            agent_id=str(self.agent.id),
            budget_id=self.budget_id,
            depth=0
        )

        # Default budget context used by spawn helpers in tests
        self.budget_ctx = BudgetContext(
            agent_id=str(self.agent.id),
            budget_id=self.budget_id,
            branch_id=self.branch_id,
            depth=0,
            max_steps=100,
            max_depth=3,
        )

    def _create_credential_secret(self, *, key: str, domain: str, value: str, name: str | None = None) -> PersistentAgentSecret:
        secret = PersistentAgentSecret(
            agent=self.agent,
            secret_type=PersistentAgentSecret.SecretType.CREDENTIAL,
            domain_pattern=domain,
            name=name or key,
            key=key,
            requested=False,
        )
        secret.set_value(value)
        secret.save()
        return secret

    def _set_browser_config(self, *, active: int | None = None, daily: int | None = None) -> None:
        config, _ = BrowserConfig.objects.get_or_create(plan_name=PlanNames.FREE)
        if active is not None:
            config.max_active_browser_tasks = active
        if daily is not None:
            config.max_browser_tasks = daily
        config.save()
        
        # Set up budget context
        self.budget_ctx = BudgetContext(
            agent_id=str(self.agent.id),
            budget_id=self.budget_id,
            branch_id=self.branch_id,
            depth=0,
            max_steps=100,
            max_depth=3
        )
    
    @patch('api.agent.tools.spawn_web_task.BrowserUseAgentTask.objects.create')
    @patch('api.models.TaskCreditService.check_and_consume_credit_for_owner', return_value={"success": True, "credit": None, "error_message": None})
    @patch('api.tasks.browser_agent_tasks.process_browser_use_task')
    def test_parallel_spawn_increments_outstanding_children(self, mock_process_task, _mock_consume_credit, mock_create):
        """Parallel spawns should each pass recursion depth=1 and increment outstanding-children counter."""
        mock_process_task.delay = MagicMock()
        task_ids = count(1)

        def _make_task(*_args, **_kwargs):
            task = MagicMock()
            task.id = next(task_ids)
            return task

        mock_create.side_effect = _make_task
        
        # Results will be stored here by each thread
        results = []
        spawn_errors = []
        observed_depths = []
        
        def spawn_task(task_num):
            """Function to be run in each thread."""
            try:
                # Set budget context for this thread
                set_budget_context(self.budget_ctx)
                
                # Small delay to ensure threads run concurrently
                time.sleep(0.01 * task_num)
                
                result = execute_spawn_web_task(
                    self.agent,
                    {"prompt": f"Task {task_num}"}
                )
                results.append(result)
                
                # Check if spawn was successful
                if result.get("status") == "error":
                    spawn_errors.append((task_num, result.get("message")))
                
                # Extract the depth that was passed to the spawned task
                if mock_process_task.delay.called:
                    calls = mock_process_task.delay.call_args_list
                    if calls:
                        # Get the most recent call's depth argument
                        last_call = calls[-1]
                        depth = last_call[1].get('depth') if last_call[1] else None
                        if depth is not None:
                            observed_depths.append((task_num, depth))
                
            except Exception as e:
                spawn_errors.append((task_num, str(e)))
        
        # Spawn 3 tasks in parallel from depth 0
        threads = []
        for i in range(3):
            thread = threading.Thread(target=spawn_task, args=(i,))
            threads.append(thread)
            thread.start()
        
        # Wait for all threads to complete
        for thread in threads:
            thread.join(timeout=5)
        
        # Check results
        self.assertEqual(len(results), 3, "All three spawns should complete")
        
        # Check for any errors (especially "recursion limit reached")
        if spawn_errors:
            self.fail(f"Some spawns failed: {spawn_errors}")
        
        # All spawned tasks should be at depth 1 since they're all spawned from depth 0 in parallel
        expected_depth = 1
        
        # Sort observed depths by task number for consistent reporting
        observed_depths.sort(key=lambda x: x[0])
        
        # Check that all tasks are at the same depth
        for task_num, depth in observed_depths:
            self.assertEqual(
                depth, expected_depth,
                f"Task {task_num} should be at depth {expected_depth}, but is at depth {depth}. "
                f"This indicates parallel spawns are incorrectly sharing/mutating the depth counter. "
                f"All observed depths: {observed_depths}"
            )
        
        # Verify the branch counter in Redis reflects all outstanding children (3)
        final_branch_depth = AgentBudgetManager.get_branch_depth(
            agent_id=str(self.agent.id),
            branch_id=self.branch_id
        )
        self.assertEqual(
            final_branch_depth, 3,
            f"Branch counter should equal number of parallel spawns; got {final_branch_depth}"
        )

    @patch('api.models.TaskCreditService.check_and_consume_credit_for_owner', return_value={"success": True, "credit": None, "error_message": None})
    @patch('api.tasks.browser_agent_tasks.process_browser_use_task.delay')
    @patch(
        'api.agent.tools.spawn_web_task.get_billing_snapshot_for_owner',
        return_value={"billing_plan": "startup", "billing_is_trial": False},
    )
    def test_spawn_web_task_stores_billing_snapshot(self, _mock_snapshot, _mock_delay, _mock_consume_credit):
        set_budget_context(self.budget_ctx)

        result = execute_spawn_web_task(self.agent, {"prompt": "Open portal.example.com"})

        self.assertEqual(result["status"], "pending")
        task = BrowserUseAgentTask.objects.get(id=result["task_id"])
        self.assertEqual(task.billing_plan, "startup")
        self.assertFalse(task.billing_is_trial)

    @override_settings(
        BROWSER_USE_TASK_EXECUTION_DISABLED=True,
        EVAL_BROWSER_TASK_SIMULATION_ENABLED=True,
    )
    @patch('api.tasks.browser_agent_tasks.process_browser_use_task')
    def test_eval_weather_browser_task_returns_synchronous_simulated_result(self, mock_process_task):
        run = EvalRun.objects.create(
            scenario_slug="monitor_pollution",
            agent=self.agent,
            initiated_by=self.user,
            status=EvalRun.Status.RUNNING,
        )
        set_budget_context(
            BudgetContext(
                agent_id=str(self.agent.id),
                budget_id=self.budget_id,
                branch_id=self.branch_id,
                depth=0,
                max_steps=100,
                max_depth=3,
                eval_run_id=str(run.id),
            )
        )

        result = execute_spawn_web_task(
            self.agent,
            {
                "prompt": (
                    "Open http://localhost:8000/eval/sim/weather/ and report the pollution index."
                ),
                "requires_vision": True,
            },
        )

        self.assertEqual(result["status"], "completed")
        self.assertEqual(result["result"]["pollution_index"], 55)
        mock_process_task.delay.assert_not_called()
        task = BrowserUseAgentTask.objects.get(id=result["task_id"])
        self.assertEqual(task.status, BrowserUseAgentTask.StatusChoices.COMPLETED)
        self.assertEqual(str(task.eval_run_id), str(run.id))
        self.assertEqual(task.steps.get().result_value["pollution_status"], "Moderate")
        self.assertEqual(
            AgentBudgetManager.get_branch_depth(
                agent_id=str(self.agent.id),
                branch_id=self.branch_id,
            ),
            0,
        )

    @patch('api.models.TaskCreditService.check_and_consume_credit_for_owner', return_value={"success": True, "credit": None, "error_message": None})
    @patch('api.tasks.browser_agent_tasks.process_browser_use_task')
    def test_sequential_spawn_depth_tracking(self, mock_process_task, _mock_consume_credit):
        """Test that sequential spawns work correctly (baseline test).
        
        This should pass even with the current buggy implementation
        since there's no parallelism.
        """
        mock_process_task.delay = MagicMock()
        
        # Set budget context
        set_budget_context(self.budget_ctx)
        
        # First spawn (from depth 0)
        result1 = execute_spawn_web_task(
            self.agent,
            {"prompt": "Task 1"}
        )
        self.assertEqual(result1.get("status"), "pending")
        
        # Get the depth of the first spawned task
        call1 = mock_process_task.delay.call_args_list[0]
        depth1 = call1[1].get('depth')
        self.assertEqual(depth1, 1, "First spawn should be at depth 1")
        
        # Reset branch depth back to 0 for second spawn
        AgentBudgetManager.set_branch_depth(
            agent_id=str(self.agent.id),
            branch_id=self.branch_id,
            depth=0
        )
        
        # Second spawn (also from depth 0)
        result2 = execute_spawn_web_task(
            self.agent,
            {"prompt": "Task 2"}
        )
        self.assertEqual(result2.get("status"), "pending")
        
        # Get the depth of the second spawned task
        call2 = mock_process_task.delay.call_args_list[1]
        depth2 = call2[1].get('depth')
        self.assertEqual(depth2, 1, "Second spawn should also be at depth 1")

    @patch('api.models.TaskCreditService.check_and_consume_credit_for_owner', return_value={"success": True, "credit": None, "error_message": None})
    @patch('api.tasks.browser_agent_tasks.process_browser_use_task')
    def test_spawn_respects_active_task_limit(self, mock_process_task, _mock_consume_credit):
        """Ensure the configured active task limit prevents new spawns."""
        mock_process_task.delay = MagicMock()

        self._set_browser_config(active=2)

        BrowserUseAgentTask.objects.create(
            agent=self.browser_agent,
            user=self.user,
            status=BrowserUseAgentTask.StatusChoices.PENDING,
            prompt="Existing task 1",
        )
        BrowserUseAgentTask.objects.create(
            agent=self.browser_agent,
            user=self.user,
            status=BrowserUseAgentTask.StatusChoices.IN_PROGRESS,
            prompt="Existing task 2",
        )

        result = execute_spawn_web_task(
            self.agent,
            {"prompt": "Should be rejected"}
        )

        self.assertEqual(result.get("status"), "error")
        self.assertIn("limit reached (2)", result.get("message", ""))
        mock_process_task.delay.assert_not_called()
        self.assertEqual(
            BrowserUseAgentTask.objects.filter(agent=self.browser_agent).count(),
            2,
            "No additional tasks should be created when the limit is reached",
        )

    @patch('util.analytics.Analytics.track_event')
    @patch('api.models.TaskCreditService.check_and_consume_credit_for_owner', return_value={"success": True, "credit": None, "error_message": None})
    @patch('api.tasks.browser_agent_tasks.process_browser_use_task')
    def test_spawn_respects_daily_task_limit(self, mock_process_task, _mock_consume_credit, mock_track_event):
        """Ensure the daily browser task limit stops additional spawns."""
        mock_process_task.delay = MagicMock()

        self._set_browser_config(daily=2)

        # Completed tasks still count toward the per-day total
        BrowserUseAgentTask.objects.create(
            agent=self.browser_agent,
            user=self.user,
            status=BrowserUseAgentTask.StatusChoices.COMPLETED,
            prompt="Daily task 1",
        )
        BrowserUseAgentTask.objects.create(
            agent=self.browser_agent,
            user=self.user,
            status=BrowserUseAgentTask.StatusChoices.FAILED,
            prompt="Daily task 2",
        )

        result = execute_spawn_web_task(
            self.agent,
            {"prompt": "Should fail due to daily limit"}
        )

        self.assertEqual(result.get("status"), "error")
        self.assertIn("Daily browser task limit reached (2)", result.get("message", ""))
        mock_process_task.delay.assert_not_called()
        mock_track_event.assert_called_once()
        event_args, event_kwargs = mock_track_event.call_args
        self.assertEqual(event_args[1], AnalyticsEvent.PERSISTENT_AGENT_BROWSER_DAILY_LIMIT_REACHED)

    @patch("api.agent.tools.spawn_web_task.logger.warning")
    @patch('api.models.TaskCreditService.check_and_consume_credit_for_owner', return_value={"success": True, "credit": None, "error_message": None})
    @patch("api.tasks.browser_agent_tasks.process_browser_use_task")
    def test_spawn_skips_invalid_stored_secret_when_not_explicitly_requested(
        self,
        mock_process_task,
        _mock_consume_credit,
        mock_logger_warning,
    ):
        mock_process_task.delay = MagicMock()
        set_budget_context(self.budget_ctx)

        self._create_credential_secret(
            key="portal_password",
            domain="https://portal.example.com",
            value="portal-secret",
        )
        invalid_secret = self._create_credential_secret(
            key="legacy_password",
            domain="https://legacy.example.com",
            value="legacy-secret",
        )
        PersistentAgentSecret.objects.filter(pk=invalid_secret.pk).update(domain_pattern="*.gov")

        result = execute_spawn_web_task(self.agent, {"prompt": "Open portal.example.com"})

        self.assertEqual(result.get("status"), "pending", result)
        mock_process_task.delay.assert_called_once()
        task = BrowserUseAgentTask.objects.get(id=result["task_id"])
        decrypted = SecretsEncryption.decrypt_secrets(task.encrypted_secrets)

        self.assertEqual(
            decrypted,
            {
                "https://portal.example.com": {
                    "portal_password": "portal-secret",
                }
            },
        )
        self.assertEqual(
            task.secret_keys,
            {"https://portal.example.com": ["portal_password"]},
        )
        self.assertNotIn("legacy_password", str(task.secret_keys))
        self.assertTrue(mock_logger_warning.called)

    @patch('api.models.TaskCreditService.check_and_consume_credit_for_owner', return_value={"success": True, "credit": None, "error_message": None})
    @patch("api.tasks.browser_agent_tasks.process_browser_use_task")
    def test_spawn_errors_for_explicit_invalid_secret(
        self,
        mock_process_task,
        _mock_consume_credit,
    ):
        mock_process_task.delay = MagicMock()
        set_budget_context(self.budget_ctx)

        secret = self._create_credential_secret(
            key="legacy_password",
            domain="https://legacy.example.com",
            value="legacy-secret",
        )
        PersistentAgentSecret.objects.filter(pk=secret.pk).update(domain_pattern="*.gov")

        result = execute_spawn_web_task(
            self.agent,
            {"prompt": "Open portal.example.com", "secrets": ["legacy_password"]},
        )

        self.assertEqual(result.get("status"), "error")
        self.assertIn("legacy_password", result.get("message", ""))
        self.assertIn("*.gov", result.get("message", ""))
        self.assertEqual(BrowserUseAgentTask.objects.count(), 0)
        mock_process_task.delay.assert_not_called()
