"""
Tests for WebSocket functionality
"""
import pytest
import json
from fastapi.testclient import TestClient
from main import app, sessions, Message


@pytest.fixture
def client():
    """Create a test client"""
    return TestClient(app)


@pytest.fixture(autouse=True)
def clear_sessions():
    """Clear sessions before each test"""
    sessions.clear()
    yield
    sessions.clear()


def create_test_session_and_character(client):
    """Helper to create a session and character"""
    session_response = client.post("/sessions/?name=TestSession")
    session_id = session_response.json()["id"]
    
    char_response = client.post(
        f"/sessions/{session_id}/characters/",
        params={"name": "TestChar", "description": "Test character"}
    )
    character_id = char_response.json()["id"]
    
    return session_id, character_id


class TestCharacterWebSocket:
    """Test character WebSocket connection"""
    
    def test_character_websocket_connection(self, client):
        """Test that character can connect to WebSocket"""
        session_id, character_id = create_test_session_and_character(client)
        
        with client.websocket_connect(f"/ws/character/{session_id}/{character_id}") as websocket:
            # Should receive history on connection
            data = websocket.receive_json()
            
            assert data["type"] == "history"
            assert "messages" in data
            assert "public_messages" in data
            assert isinstance(data["messages"], list)
            assert isinstance(data["public_messages"], list)
    
    def test_character_websocket_invalid_session(self, client):
        """Test connection with invalid session ID"""
        with pytest.raises(Exception):
            with client.websocket_connect(f"/ws/character/fake-session/fake-char"):
                pass
    
    def test_character_websocket_invalid_character(self, client):
        """Test connection with invalid character ID"""
        session_response = client.post("/sessions/?name=TestSession")
        session_id = session_response.json()["id"]
        
        with pytest.raises(Exception):
            with client.websocket_connect(f"/ws/character/{session_id}/fake-character"):
                pass
    
    def test_character_receives_history(self, client):
        """Test that character receives their conversation history"""
        session_id, character_id = create_test_session_and_character(client)
        
        # Add a message to character's history manually
        session = sessions[session_id]
        character = session.characters[character_id]
        test_message = Message(sender="storyteller", content="Welcome!")
        character.conversation_history.append(test_message)
        
        with client.websocket_connect(f"/ws/character/{session_id}/{character_id}") as websocket:
            data = websocket.receive_json()
            
            assert data["type"] == "history"
            assert len(data["messages"]) == 1
            assert data["messages"][0]["content"] == "Welcome!"
    
    def test_character_sends_message(self, client):
        """Test character sending a message"""
        session_id, character_id = create_test_session_and_character(client)
        
        with client.websocket_connect(f"/ws/character/{session_id}/{character_id}") as websocket:
            # Receive initial history
            websocket.receive_json()
            
            # Send a message
            websocket.send_json({
                "type": "message",
                "content": "I search the room",
                "visibility": "private"
            })
            
            # Verify message was added to character's history
            session = sessions[session_id]
            character = session.characters[character_id]
            
            assert len(character.conversation_history) > 0
            assert character.pending_response is True


class TestStorytellerWebSocket:
    """Test storyteller WebSocket connection"""
    
    def test_storyteller_websocket_connection(self, client):
        """Test that storyteller can connect to WebSocket"""
        session_response = client.post("/sessions/?name=TestSession")
        session_id = session_response.json()["id"]
        
        with client.websocket_connect(f"/ws/storyteller/{session_id}") as websocket:
            data = websocket.receive_json()
            
            assert data["type"] == "session_state"
            assert "characters" in data
            assert "current_scene" in data
            assert "public_messages" in data
    
    def test_storyteller_sees_all_characters(self, client):
        """Test that storyteller receives all character data"""
        session_response = client.post("/sessions/?name=TestSession")
        session_id = session_response.json()["id"]
        
        # Add two characters
        char1 = client.post(
            f"/sessions/{session_id}/characters/",
            params={"name": "Char1", "description": "First"}
        ).json()
        
        char2 = client.post(
            f"/sessions/{session_id}/characters/",
            params={"name": "Char2", "description": "Second"}
        ).json()
        
        with client.websocket_connect(f"/ws/storyteller/{session_id}") as websocket:
            data = websocket.receive_json()
            
            assert len(data["characters"]) == 2
            assert char1["id"] in data["characters"]
            assert char2["id"] in data["characters"]
    
    def test_storyteller_websocket_invalid_session(self, client):
        """Test storyteller connection with invalid session"""
        with pytest.raises(Exception):
            with client.websocket_connect(f"/ws/storyteller/fake-session"):
                pass


class TestMessageRouting:
    """Test message routing between character and storyteller"""
    
    def test_private_message_routing(self, client):
        """Test that private messages route to storyteller only"""
        session_id, character_id = create_test_session_and_character(client)
        
        # Connect character
        with client.websocket_connect(f"/ws/character/{session_id}/{character_id}") as char_ws:
            # Receive initial history
            char_ws.receive_json()
            
            # Send private message
            char_ws.send_json({
                "type": "message",
                "content": "Secret action",
                "visibility": "private"
            })
            
            # Verify it's in character's private history
            session = sessions[session_id]
            character = session.characters[character_id]
            
            assert len(character.conversation_history) == 1
            assert character.conversation_history[0].visibility == "private"
            
            # Verify it's NOT in public messages
            assert len(session.public_messages) == 0
    
    def test_public_message_routing(self, client):
        """Test that public messages are added to public feed"""
        session_id, character_id = create_test_session_and_character(client)
        
        with client.websocket_connect(f"/ws/character/{session_id}/{character_id}") as char_ws:
            # Receive initial history
            char_ws.receive_json()
            
            # Send public message
            char_ws.send_json({
                "type": "message",
                "content": "I wave to everyone",
                "visibility": "public"
            })
            
            # Verify it's in public messages
            session = sessions[session_id]
            assert len(session.public_messages) == 1
            assert session.public_messages[0].visibility == "public"
            assert session.public_messages[0].content == "I wave to everyone"
    
    def test_mixed_message_routing(self, client):
        """Test that mixed messages go to both feeds"""
        session_id, character_id = create_test_session_and_character(client)
        
        with client.websocket_connect(f"/ws/character/{session_id}/{character_id}") as char_ws:
            # Receive initial history
            char_ws.receive_json()
            
            # Send mixed message
            char_ws.send_json({
                "type": "message",
                "content": "Combined message",
                "visibility": "mixed",
                "public_content": "I greet the guard",
                "private_content": "I scan for weaknesses"
            })
            
            session = sessions[session_id]
            character = session.characters[character_id]
            
            # Should be in public feed
            assert len(session.public_messages) == 1
            # Should be in private history
            assert len(character.conversation_history) == 1
            # Should have pending response
            assert character.pending_response is True


class TestStorytellerResponses:
    """Test storyteller responding to characters"""
    
    def test_storyteller_responds_to_character(self, client):
        """Test storyteller sending response to a character"""
        session_id, character_id = create_test_session_and_character(client)
        
        # Add a pending message from character
        session = sessions[session_id]
        character = session.characters[character_id]
        character.conversation_history.append(
            Message(sender="character", content="What do I see?")
        )
        character.pending_response = True
        
        # Connect storyteller
        with client.websocket_connect(f"/ws/storyteller/{session_id}") as st_ws:
            # Receive initial state
            st_ws.receive_json()
            
            # Send response
            st_ws.send_json({
                "type": "respond_to_character",
                "character_id": character_id,
                "content": "You see a vast chamber"
            })
            
            # Verify response added to character's history
            assert len(character.conversation_history) == 2
            assert character.conversation_history[1].sender == "storyteller"
            assert character.conversation_history[1].content == "You see a vast chamber"
            # Pending flag should be cleared
            assert character.pending_response is False


class TestSceneNarration:
    """Test scene narration broadcasting"""
    
    def test_storyteller_narrates_scene(self, client):
        """Test storyteller narrating a scene"""
        session_id, character_id = create_test_session_and_character(client)
        
        with client.websocket_connect(f"/ws/storyteller/{session_id}") as st_ws:
            # Receive initial state
            st_ws.receive_json()
            
            # Narrate scene
            scene_text = "You enter a dark cavern filled with ancient relics"
            st_ws.send_json({
                "type": "narrate_scene",
                "content": scene_text
            })
            
            # Verify scene updated
            session = sessions[session_id]
            assert session.current_scene == scene_text
            assert scene_text in session.scene_history


class TestConnectionManager:
    """Test connection management"""
    
    def test_multiple_character_connections(self, client):
        """Test multiple characters can connect simultaneously"""
        session_response = client.post("/sessions/?name=TestSession")
        session_id = session_response.json()["id"]
        
        # Create two characters
        char1_response = client.post(
            f"/sessions/{session_id}/characters/",
            params={"name": "Char1", "description": "First"}
        )
        char1_id = char1_response.json()["id"]
        
        char2_response = client.post(
            f"/sessions/{session_id}/characters/",
            params={"name": "Char2", "description": "Second"}
        )
        char2_id = char2_response.json()["id"]
        
        # Connect both
        with client.websocket_connect(f"/ws/character/{session_id}/{char1_id}") as ws1:
            with client.websocket_connect(f"/ws/character/{session_id}/{char2_id}") as ws2:
                # Both should receive history
                data1 = ws1.receive_json()
                data2 = ws2.receive_json()
                
                assert data1["type"] == "history"
                assert data2["type"] == "history"
    
    def test_storyteller_and_character_simultaneous(self, client):
        """Test storyteller and character can be connected at same time"""
        session_id, character_id = create_test_session_and_character(client)
        
        with client.websocket_connect(f"/ws/storyteller/{session_id}") as st_ws:
            with client.websocket_connect(f"/ws/character/{session_id}/{character_id}") as char_ws:
                # Both should connect successfully
                st_data = st_ws.receive_json()
                char_data = char_ws.receive_json()
                
                assert st_data["type"] == "session_state"
                assert char_data["type"] == "history"


class TestMessagePersistence:
    """Test that messages persist in session"""
    
    def test_messages_persist_after_disconnect(self, client):
        """Test that messages remain after WebSocket disconnect"""
        session_id, character_id = create_test_session_and_character(client)
        
        # Connect and send message
        with client.websocket_connect(f"/ws/character/{session_id}/{character_id}") as websocket:
            websocket.receive_json()
            websocket.send_json({
                "type": "message",
                "content": "Test message",
                "visibility": "private"
            })
        
        # WebSocket disconnected, check if message persists
        session = sessions[session_id]
        character = session.characters[character_id]
        
        assert len(character.conversation_history) == 1
        assert character.conversation_history[0].content == "Test message"
    
    def test_reconnect_receives_history(self, client):
        """Test that reconnecting receives previous messages"""
        session_id, character_id = create_test_session_and_character(client)
        
        # First connection - send message
        with client.websocket_connect(f"/ws/character/{session_id}/{character_id}") as websocket:
            websocket.receive_json()
            websocket.send_json({
                "type": "message",
                "content": "First message",
                "visibility": "private"
            })
        
        # Second connection - should receive history
        with client.websocket_connect(f"/ws/character/{session_id}/{character_id}") as websocket:
            data = websocket.receive_json()
            
            assert data["type"] == "history"
            assert len(data["messages"]) == 1
            assert data["messages"][0]["content"] == "First message"