"""
Tests for agentId tracking on agent messages.

Verifies:
1. Agent messages are saved with the correct agentId.
2. User (visitor) messages have no agentId (null/absent).
3. When Agent A handles part of a conversation, it times out, and Agent B
   claims the re-ring — messages from each agent show their own correct,
   DIFFERENT agentId.
4. History retrieval endpoints return the agentId field.
"""
import pytest
from app.database import get_database
from app.repositories.session_repository import create_session
from app.repositories.conversation_repository import get_or_create_conversation
from app.repositories.message_repository import create_message, get_messages_by_conversation_id
from app.services.presence import add_active_session


@pytest.mark.anyio
async def test_agent_message_stores_agent_id(client):
    """An agent message saved with agent_id should have agentId in the document."""
    db = get_database()
    sess = await create_session(db, site_id="site_demo_123")
    sid = sess["sessionId"]
    conv = await get_or_create_conversation(db, sid, "site_demo_123")
    cid = conv["conversationId"]

    msg_doc = await create_message(
        db,
        conversation_id=cid,
        session_id=sid,
        role="agent",
        content="Hello from agent_alpha",
        agent_id="agent_alpha"
    )

    assert msg_doc["agentId"] == "agent_alpha"
    assert msg_doc["role"] == "agent"

    # Verify persisted in DB
    messages = await get_messages_by_conversation_id(db, cid)
    agent_msgs = [m for m in messages if m["role"] == "agent"]
    assert len(agent_msgs) == 1
    assert agent_msgs[0]["agentId"] == "agent_alpha"


@pytest.mark.anyio
async def test_user_message_has_no_agent_id(client):
    """A user (visitor) message should have no agentId field."""
    db = get_database()
    sess = await create_session(db, site_id="site_demo_123")
    sid = sess["sessionId"]
    conv = await get_or_create_conversation(db, sid, "site_demo_123")
    cid = conv["conversationId"]

    msg_doc = await create_message(
        db,
        conversation_id=cid,
        session_id=sid,
        role="user",
        content="Hello from visitor"
    )

    assert "agentId" not in msg_doc
    assert msg_doc["role"] == "user"

    # Verify persisted in DB
    messages = await get_messages_by_conversation_id(db, cid)
    user_msgs = [m for m in messages if m["role"] == "user"]
    assert len(user_msgs) == 1
    assert "agentId" not in user_msgs[0]


@pytest.mark.anyio
async def test_rering_different_agents_have_different_agent_ids(client):
    """
    Scenario: Agent A handles part of a conversation, it times out / disconnects,
    and Agent B claims the re-ring. Messages from each agent should record their
    own correct, DIFFERENT agentId.
    """
    db = get_database()
    sess = await create_session(db, site_id="site_demo_123")
    sid = sess["sessionId"]
    conv = await get_or_create_conversation(db, sid, "site_demo_123")
    cid = conv["conversationId"]

    # Agent A sends two messages
    await create_message(db, cid, sid, "agent", "Agent A message 1", agent_id="agent_alpha")
    await create_message(db, cid, sid, "agent", "Agent A message 2", agent_id="agent_alpha")

    # Visitor sends a message (no agentId)
    await create_message(db, cid, sid, "user", "Visitor reply")

    # Agent B claims the re-ring and sends messages
    await create_message(db, cid, sid, "agent", "Agent B message 1", agent_id="agent_beta")
    await create_message(db, cid, sid, "agent", "Agent B message 2", agent_id="agent_beta")

    # Retrieve all messages
    messages = await get_messages_by_conversation_id(db, cid)
    assert len(messages) == 5

    # Verify Agent A's messages
    assert messages[0]["agentId"] == "agent_alpha"
    assert messages[1]["agentId"] == "agent_alpha"

    # Verify visitor message has no agentId
    assert "agentId" not in messages[2]

    # Verify Agent B's messages
    assert messages[3]["agentId"] == "agent_beta"
    assert messages[4]["agentId"] == "agent_beta"

    # Ensure the two agents are actually different
    agent_ids = set(m.get("agentId") for m in messages if m.get("agentId"))
    assert agent_ids == {"agent_alpha", "agent_beta"}


@pytest.mark.anyio
async def test_handoff_message_endpoint_stores_agent_id(client):
    """
    The POST /api/v1/handoff/message endpoint should store the agent's ID
    on the message document and return it in the session history.
    """
    db = get_database()
    sess = await create_session(db, site_id="site_demo_123")
    sid = sess["sessionId"]

    # Set up in-memory session as IN_CHAT with agent_alpha
    from unittest.mock import AsyncMock
    mock_ws = AsyncMock()
    add_active_session(sid, "site_demo_123", mock_ws, status_val="IN_CHAT", assigned_agent_id="agent_alpha")

    # Send agent message via the handoff endpoint
    resp = client.post("/api/v1/handoff/message", json={
        "sessionId": sid,
        "message": "Hello from agent alpha",
        "agentId": "agent_alpha"
    })
    assert resp.status_code == 200

    # Retrieve session history and verify agentId is present
    history_resp = client.get(f"/api/v1/session/{sid}/history?siteId=site_demo_123")
    assert history_resp.status_code == 200
    history = history_resp.json()

    agent_msgs = [m for m in history["messages"] if m["role"] == "agent"]
    assert len(agent_msgs) >= 1
    assert agent_msgs[0]["agentId"] == "agent_alpha"


@pytest.mark.anyio
async def test_visitor_message_has_null_agent_id_in_history(client):
    """
    User/visitor messages should have agentId=null in the session history response.
    """
    db = get_database()
    sess = await create_session(db, site_id="site_demo_123")
    sid = sess["sessionId"]

    # Set up in-memory session
    from unittest.mock import AsyncMock
    mock_ws = AsyncMock()
    add_active_session(sid, "site_demo_123", mock_ws, status_val="IN_CHAT", assigned_agent_id="agent_alpha")

    # Send a visitor message
    resp = client.post("/api/v1/message", json={
        "sessionId": sid,
        "message": "Hello from visitor"
    })
    assert resp.status_code == 200

    # Retrieve session history and verify visitor messages have no agentId
    history_resp = client.get(f"/api/v1/session/{sid}/history?siteId=site_demo_123")
    assert history_resp.status_code == 200
    history = history_resp.json()

    user_msgs = [m for m in history["messages"] if m["role"] == "user"]
    assert len(user_msgs) >= 1
    assert user_msgs[0].get("agentId") is None


@pytest.mark.anyio
async def test_rering_scenario_via_endpoints(client):
    """
    End-to-end re-ring scenario via HTTP endpoints:
    1. Agent A claims session and sends a message.
    2. Session goes to ENDED (simulated).
    3. Session re-rings, Agent B claims and sends a message.
    4. Verify Agent A's message has agentId=agent_alpha (in conv 1).
    5. Verify Agent B's message has agentId=agent_beta (in conv 2).
    
    Note: After end+re-ring, get_or_create_conversation creates a new
    conversation. So the two agents' messages end up in separate conversations.
    This test verifies both independently.
    """
    db = get_database()
    sess = await create_session(db, site_id="site_demo_123")
    sid = sess["sessionId"]

    from unittest.mock import AsyncMock
    mock_ws = AsyncMock()
    add_active_session(sid, "site_demo_123", mock_ws, status_val="RINGING")

    # Agent A claims
    accept_resp_a = client.post("/api/v1/handoff/accept", json={
        "sessionId": sid,
        "agentId": "agent_alpha"
    })
    assert accept_resp_a.status_code == 200
    assert accept_resp_a.json()["assignedAgentId"] == "agent_alpha"

    # Agent A sends a message
    msg_resp_a = client.post("/api/v1/handoff/message", json={
        "sessionId": sid,
        "message": "Agent Alpha here",
        "agentId": "agent_alpha"
    })
    assert msg_resp_a.status_code == 200

    # Verify Agent A's message has correct agentId via direct DB query
    conv_a = await db["conversations"].find_one({"sessionId": sid, "status": {"$ne": "ended"}})
    assert conv_a is not None
    conv_a_id = conv_a["conversationId"]
    msgs_a = await get_messages_by_conversation_id(db, conv_a_id)
    agent_msgs_a = [m for m in msgs_a if m["role"] == "agent"]
    assert len(agent_msgs_a) == 1
    assert agent_msgs_a[0]["agentId"] == "agent_alpha"

    # End chat
    end_resp = client.post("/api/v1/handoff/end", json={
        "sessionId": sid,
        "endedBy": "agent"
    })
    assert end_resp.status_code == 200

    # Re-ring: reset session status to RINGING and clear assignment
    await db["sessions"].update_one(
        {"sessionId": sid},
        {"$set": {"status": "RINGING", "assignedAgentId": None, "assignedAt": None, "mode": "rule"}}
    )
    add_active_session(sid, "site_demo_123", mock_ws, status_val="RINGING")

    # Agent B claims
    accept_resp_b = client.post("/api/v1/handoff/accept", json={
        "sessionId": sid,
        "agentId": "agent_beta"
    })
    assert accept_resp_b.status_code == 200
    assert accept_resp_b.json()["assignedAgentId"] == "agent_beta"

    # Agent B sends a message
    msg_resp_b = client.post("/api/v1/handoff/message", json={
        "sessionId": sid,
        "message": "Agent Beta here",
        "agentId": "agent_beta"
    })
    assert msg_resp_b.status_code == 200

    # Verify Agent B's message has correct agentId via direct DB query
    # Find the newest conversation (not the ended one)
    conv_b = await db["conversations"].find_one(
        {"sessionId": sid, "conversationId": {"$ne": conv_a_id}}
    )
    assert conv_b is not None, "Expected a second conversation after re-ring"
    conv_b_id = conv_b["conversationId"]
    msgs_b = await get_messages_by_conversation_id(db, conv_b_id)
    agent_msgs_b = [m for m in msgs_b if m["role"] == "agent"]
    assert len(agent_msgs_b) == 1
    assert agent_msgs_b[0]["agentId"] == "agent_beta"

    # Crucially: the two agents have DIFFERENT agentIds
    assert agent_msgs_a[0]["agentId"] != agent_msgs_b[0]["agentId"]


def test_agent_message_via_agent_websocket(client):
    """
    When an agent sends a message over /ws/agent WebSocket,
    the persisted MongoDB message document must store the sending agent's agentId.
    """
    import asyncio
    db = get_database()
    sess = asyncio.run(create_session(db, site_id="site_demo_123"))
    sid = sess["sessionId"]

    from unittest.mock import AsyncMock
    mock_ws = AsyncMock()
    add_active_session(sid, "site_demo_123", mock_ws, status_val="IN_CHAT", assigned_agent_id="agent_gamma")

    # Connect to agent WebSocket and send an agent message
    with client.websocket_connect("/ws/agent?agentId=agent_gamma") as ws:
        ws.send_json({
            "type": "agent_message",
            "sessionId": sid,
            "message": "Live reply via agent websocket"
        })

    # Verify message in DB has agentId = agent_gamma
    conv = asyncio.run(db["conversations"].find_one({"sessionId": sid}))
    assert conv is not None
    msgs = asyncio.run(get_messages_by_conversation_id(db, conv["conversationId"]))
    agent_msgs = [m for m in msgs if m["role"] == "agent"]
    assert len(agent_msgs) == 1
    assert agent_msgs[0]["agentId"] == "agent_gamma"
    assert agent_msgs[0]["content"] == "Live reply via agent websocket"