import pytest

# Test prompts and data
CONTEXT_PROMPT = "Who is the CEO of Microsoft?"
FOLLOW_UP_PROMPT = "When was he born?"
RAG_DOC_TITLE = "Fictional Company History"
RAG_DOC_TEXT = "The company AlphaCorp was founded in 2021 by Jane Doe. Their primary product is a smart home device called 'Nexus'."
RAG_PROMPT = "Who founded AlphaCorp and what is their main product?"

@pytest.mark.asyncio
async def test_chat_in_session_lifecycle(http_client):
    """
    Tests a full session lifecycle from creation to conversational memory.
    This test is a single, sequential unit.
    """
    print("\n--- Running test_chat_in_session_lifecycle ---")
    
    # 1. Create a new session
    payload = {"user_id": "integration_tester_lifecycle", "model": "deepseek"}
    response = await http_client.post("/sessions", json=payload)
    assert response.status_code == 200
    session_id = response.json()["id"]
    print(f"✅ Session created successfully with ID: {session_id}")
    
    # 2. First chat turn to establish context
    chat_payload_1 = {"prompt": CONTEXT_PROMPT}
    response_1 = await http_client.post(f"/sessions/{session_id}/chat", json=chat_payload_1)
    assert response_1.status_code == 200
    assert "Satya Nadella" in response_1.json()["answer"]
    assert response_1.json()["model_used"] == "deepseek"
    print("✅ Chat Turn 1 (context) test passed.")
    
    # 3. Second chat turn (follow-up) to test conversational memory
    chat_payload_2 = {"prompt": FOLLOW_UP_PROMPT}
    response_2 = await http_client.post(f"/sessions/{session_id}/chat", json=chat_payload_2)
    assert response_2.status_code == 200
    assert "1967" in response_2.json()["answer"]
    assert response_2.json()["model_used"] == "deepseek"
    print("✅ Chat Turn 2 (follow-up) test passed.")
    
    # 4. Cleanup (optional, but good practice if not using a test database that resets)
    # The session data would typically be cleaned up by the database teardown.

@pytest.mark.asyncio
async def test_chat_with_model_switch(http_client, session_id):
    """Tests switching models within an existing session."""
    print("\n--- Running test_chat_with_model_switch ---")
    
    # Send a message to the new session with a different model
    payload_gemini = {"prompt": "What is the capital of France?", "model": "gemini"}
    response_gemini = await http_client.post(f"/sessions/{session_id}/chat", json=payload_gemini)
    assert response_gemini.status_code == 200
    assert "Paris" in response_gemini.json()["answer"]
    assert response_gemini.json()["model_used"] == "gemini"
    print("✅ Chat (Model Switch to Gemini) test passed.")

    # Switch back to the original model
    payload_deepseek = {"prompt": "What is the largest ocean?", "model": "deepseek"}
    response_deepseek = await http_client.post(f"/sessions/{session_id}/chat", json=payload_deepseek)
    assert response_deepseek.status_code == 200
    assert "Pacific Ocean" in response_deepseek.json()["answer"]
    assert response_deepseek.json()["model_used"] == "deepseek"
    print("✅ Chat (Model Switch back to DeepSeek) test passed.")

@pytest.mark.asyncio
async def test_chat_with_document_retrieval(http_client):
    """
    Tests injecting a document and using it for retrieval-augmented generation.
    This test creates its own session and document for isolation.
    """
    print("\n--- Running test_chat_with_document_retrieval ---")
    
    # Create a new session for this RAG test
    session_response = await http_client.post("/sessions", json={"user_id": "rag_tester", "model": "deepseek"})
    assert session_response.status_code == 200
    rag_session_id = session_response.json()["id"]
    
    # Add a new document with specific content for retrieval
    doc_data = {"title": RAG_DOC_TITLE, "text": RAG_DOC_TEXT}
    add_doc_response = await http_client.post("/documents", json=doc_data)
    assert add_doc_response.status_code == 200
    try:
        message = add_doc_response.json().get("message", "")
        rag_document_id = int(message.split(" with ID ")[-1])
        print(f"Document for RAG created with ID: {rag_document_id}")
    except (ValueError, IndexError):
        pytest.fail("Could not parse document ID from response message.")
    
    try:
        chat_payload = {
            "prompt": RAG_PROMPT,
            "document_id": rag_document_id,
            "model": "deepseek",
            "load_faiss_retriever": True
        }
        chat_response = await http_client.post(f"/sessions/{rag_session_id}/chat", json=chat_payload)
        
        assert chat_response.status_code == 200
        chat_data = chat_response.json()
        assert "Jane Doe" in chat_data["answer"]
        assert "Nexus" in chat_data["answer"]
        print("✅ Chat with document retrieval test passed.")
    finally:
        # Clean up the document after the test
        delete_response = await http_client.delete(f"/documents/{rag_document_id}")
        assert delete_response.status_code == 200
        print(f"Document {rag_document_id} deleted successfully.")