import pytest
import httpx

# The base URL for the local server
BASE_URL = "http://127.0.0.1:8000"
# Use a specific, context-setting prompt for the conversational test
CONTEXT_PROMPT = "Who is the CEO of Microsoft?"
FOLLOW_UP_PROMPT = "When was he born?"

# Global variables to pass state between sequential tests
created_document_id = None
created_session_id = None

async def test_root_endpoint():
    """Tests if the root endpoint is alive."""
    print("\n--- Running test_root_endpoint ---")
    async with httpx.AsyncClient() as client:
        response = await client.get(f"{BASE_URL}/")
        assert response.status_code == 200
        assert response.json() == {"status": "AI Model Hub is running!"}
    print("✅ Root endpoint test passed.")

# --- Session and Chat Lifecycle Tests ---

async def test_create_session():
    """Tests creating a new chat session and saves the ID for the next test."""
    global created_session_id
    print("\n--- Running test_create_session ---")
    url = f"{BASE_URL}/sessions"
    payload = {"user_id": "integration_tester", "model": "deepseek"}
    
    async with httpx.AsyncClient() as client:
        response = await client.post(url, json=payload)

    assert response.status_code == 200, f"Failed to create session. Response: {response.text}"
    response_data = response.json()
    assert "id" in response_data
    created_session_id = response_data["id"]
    print(f"✅ Session created successfully with ID: {created_session_id}")

async def test_chat_in_session_turn_1():
    """Tests sending the first message to establish context."""
    print("\n--- Running test_chat_in_session (Turn 1) ---")
    assert created_session_id is not None, "Session ID was not set."
    
    url = f"{BASE_URL}/sessions/{created_session_id}/chat"
    payload = {"prompt": CONTEXT_PROMPT}
    
    async with httpx.AsyncClient(timeout=60.0) as client:
        response = await client.post(url, json=payload)
        
    assert response.status_code == 200, f"Chat request failed. Response: {response.text}"
    response_data = response.json()
    # Check that the answer mentions the CEO's name
    assert "Satya Nadella" in response_data["answer"]
    print("✅ Chat Turn 1 (context) test passed.")

async def test_chat_in_session_turn_2_follow_up():
    """
    Tests sending a follow-up question to verify conversational memory.
    """
    print("\n--- Running test_chat_in_session (Turn 2 - Follow-up) ---")
    assert created_session_id is not None, "Session ID was not set."
    
    url = f"{BASE_URL}/sessions/{created_session_id}/chat"
    payload = {"prompt": FOLLOW_UP_PROMPT}
    
    async with httpx.AsyncClient(timeout=60.0) as client:
        response = await client.post(url, json=payload)
        
    assert response.status_code == 200, f"Follow-up chat request failed. Response: {response.text}"
    response_data = response.json()
    # Check that the answer contains the birth year, proving it understood "he"
    assert "1967" in response_data["answer"]
    print("✅ Chat Turn 2 (follow-up) test passed.")

async def test_get_session_history():
    """Tests retrieving the full message history for the session."""
    print("\n--- Running test_get_session_history ---")
    assert created_session_id is not None, "Session ID was not set."
    
    url = f"{BASE_URL}/sessions/{created_session_id}/messages"
    async with httpx.AsyncClient() as client:
        response = await client.get(url)
        
    assert response.status_code == 200
    response_data = response.json()
    
    assert response_data["session_id"] == created_session_id
    # After two turns, there should be 4 messages (2 user, 2 assistant)
    assert len(response_data["messages"]) >= 4
    assert response_data["messages"][0]["content"] == CONTEXT_PROMPT
    assert response_data["messages"][2]["content"] == FOLLOW_UP_PROMPT
    print("✅ Get session history test passed.")

# --- Document Management Lifecycle Tests ---
# (These tests remain unchanged)
async def test_add_document_for_lifecycle():
    global created_document_id
    print("\n--- Running test_add_document (for lifecycle) ---")
    url = f"{BASE_URL}/documents"
    doc_data = {"title": "Lifecycle Test Doc", "text": "This doc will be listed and deleted."}
    
    async with httpx.AsyncClient(timeout=30.0) as client:
        response = await client.post(url, json=doc_data)

    assert response.status_code == 200
    try:
        message = response.json().get("message", "")
        created_document_id = int(message.split(" with ID ")[-1])
    except (ValueError, IndexError):
        pytest.fail("Could not parse document ID from response message.")
    print(f"✅ Document for lifecycle test created with ID: {created_document_id}")

async def test_list_documents():
    print("\n--- Running test_list_documents ---")
    assert created_document_id is not None, "Document ID was not set."
    
    url = f"{BASE_URL}/documents"
    async with httpx.AsyncClient() as client:
        response = await client.get(url)

    assert response.status_code == 200
    ids_in_response = {doc["id"] for doc in response.json()["documents"]}
    assert created_document_id in ids_in_response
    print("✅ Document list test passed.")

async def test_delete_document():
    print("\n--- Running test_delete_document ---")
    assert created_document_id is not None, "Document ID was not set."

    url = f"{BASE_URL}/documents/{created_document_id}"
    async with httpx.AsyncClient() as client:
        response = await client.delete(url)
    
    assert response.status_code == 200
    assert response.json()["document_id"] == created_document_id
    print("✅ Document delete test passed.")