import pytest
from unittest.mock import MagicMock, patch
import multi_agent_conductor
from models import Ticket, WorkerContext
import events
import asyncio
import concurrent.futures

class MockDialog:
 def __init__(self, approved, final_payload=None):
  self.approved = approved
  self.final_payload = final_payload

 def wait(self):
 # Match the new return format: a dictionary
  res = {'approved': self.approved, 'abort': False}
  if self.final_payload:
   res.update(self.final_payload)
  return res

@pytest.fixture
def mock_ai_client() -> None:
 with patch("ai_client.send") as mock_send:
  mock_send.return_value = "Task completed"
  yield mock_send

@pytest.mark.asyncio
async def test_confirm_spawn_pushed_to_queue():
 event_queue = events.AsyncEventQueue()
 ticket_id = "T1"
 role = "Tier 3 Worker"
 prompt = "Original Prompt"
 context_md = "Original Context"
 # Start confirm_spawn in a thread since it blocks with time.sleep

 def run_confirm():
  return multi_agent_conductor.confirm_spawn(role, prompt, context_md, event_queue, ticket_id)
 loop = asyncio.get_running_loop()
 with concurrent.futures.ThreadPoolExecutor() as executor:
  future = loop.run_in_executor(executor, run_confirm)
  # Wait for the event to appear in the queue
  event_name, payload = await event_queue.get()
  assert event_name == "mma_spawn_approval"
  assert payload["ticket_id"] == ticket_id
  assert payload["role"] == role
  assert payload["prompt"] == prompt
  assert payload["context_md"] == context_md
  assert "dialog_container" in payload
  # Simulate GUI injecting a dialog
  payload["dialog_container"][0] = MockDialog(True, {"prompt": "Modified Prompt", "context_md": "Modified Context"})
  approved, final_prompt, final_context = await future
  assert approved is True
  assert final_prompt == "Modified Prompt"
  assert final_context == "Modified Context"

@patch("multi_agent_conductor.confirm_spawn")
def test_run_worker_lifecycle_approved(mock_confirm, mock_ai_client):
 ticket = Ticket(id="T1", description="desc", status="todo", assigned_to="user")
 context = WorkerContext(ticket_id="T1", model_name="model", messages=[])
 event_queue = events.AsyncEventQueue()
 mock_confirm.return_value = (True, "Modified Prompt", "Modified Context")
 multi_agent_conductor.run_worker_lifecycle(ticket, context, event_queue=event_queue)
 mock_confirm.assert_called_once()
 # Check that ai_client.send was called with modified values
 args, kwargs = mock_ai_client.call_args
 assert kwargs["user_message"] == "Modified Prompt"
 assert kwargs["md_content"] == "Modified Context"
 assert ticket.status == "completed"

@patch("multi_agent_conductor.confirm_spawn")
def test_run_worker_lifecycle_rejected(mock_confirm, mock_ai_client):
 ticket = Ticket(id="T1", description="desc", status="todo", assigned_to="user")
 context = WorkerContext(ticket_id="T1", model_name="model", messages=[])
 event_queue = events.AsyncEventQueue()
 mock_confirm.return_value = (False, "Original Prompt", "Original Context")
 result = multi_agent_conductor.run_worker_lifecycle(ticket, context, event_queue=event_queue)
 mock_confirm.assert_called_once()
 mock_ai_client.assert_not_called()
 assert ticket.status == "blocked"
 assert "Spawn rejected by user" in ticket.blocked_reason
 assert "BLOCKED" in result