import pytest
import httpx

# The base URL for the local server
BASE_URL = "http://127.0.0.1:8000"
TEST_PROMPT = "Explain the theory of relativity in one sentence."

# This global variable will store the ID of a document created by one test
# so that it can be used by subsequent tests for listing and deletion.
created_document_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.")

# --- Chat Endpoint Tests ---

async def test_chat_endpoint_deepseek():
    """Tests the /chat endpoint using the 'deepseek' model."""
    print("\n--- Running test_chat_endpoint_deepseek ---")
    url = f"{BASE_URL}/chat"
    payload = {"prompt": TEST_PROMPT, "model": "deepseek"}
    async with httpx.AsyncClient(timeout=60.0) as client:
        response = await client.post(url, json=payload)
    assert response.status_code == 200
    data = response.json()
    assert "answer" in data
    assert data["model_used"] == "deepseek"
    print("✅ DeepSeek chat test passed.")

async def test_chat_endpoint_gemini():
    """Tests the /chat endpoint using the 'gemini' model."""
    print("\n--- Running test_chat_endpoint_gemini ---")
    url = f"{BASE_URL}/chat"
    payload = {"prompt": TEST_PROMPT, "model": "gemini"}
    async with httpx.AsyncClient(timeout=60.0) as client:
        response = await client.post(url, json=payload)
    assert response.status_code == 200
    data = response.json()
    assert "answer" in data
    assert data["model_used"] == "gemini"
    print("✅ Gemini chat test passed.")

# --- Document Management Lifecycle Tests ---

async def test_add_document_for_lifecycle():
    """
    Adds a document and saves its ID to be used by the list and delete tests.
    """
    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, f"Failed to add document. Response: {response.text}"
    response_data = response.json()
    message = response_data.get("message", "")
    assert "added successfully with ID" in message

    # Extract the ID from the success message for the next tests
    try:
        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():
    """
    Tests the GET /documents endpoint to ensure it returns a list
    that includes the document created in the previous test.
    """
    print("\n--- Running test_list_documents ---")
    assert created_document_id is not None, "Document ID was not set by the add test."
    
    url = f"{BASE_URL}/documents"
    async with httpx.AsyncClient() as client:
        response = await client.get(url)

    assert response.status_code == 200
    response_data = response.json()
    assert "documents" in response_data
    
    # Check if the list of documents contains the one we just created
    ids_in_response = {doc["id"] for doc in response_data["documents"]}
    assert created_document_id in ids_in_response, f"Document ID {created_document_id} not found in list."
    print("✅ Document list test passed.")

async def test_delete_document():
    """
    Tests the DELETE /documents/{id} endpoint to remove the document
    created at the start of the lifecycle tests.
    """
    print("\n--- Running test_delete_document ---")
    assert created_document_id is not None, "Document ID was not set by the add test."

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