"""
Integration tests for Docling Serve SDK — runs against a LIVE server.

Prerequisites
-------------
1. Start a Docling Serve instance:
       docker run -p 5001:5001 ds4sd/docling-serve
   Or any other method you use to run the server.

2. Place a small test PDF in tests/fixtures/sample.pdf
   (a 1–2 page PDF is ideal for fast tests).

3. Run the tests:
       pytest tests/test_integration.py -v --tb=short

   Optionally override the server URL or API key via env vars:
       DOCLING_SERVER_URL=http://localhost:5001
       DOCLING_API_KEY=your-key-here

   To skip async tests (if event-loop issues):
       pytest tests/test_integration.py -v -k "not _async"

All tests are marked with ``@pytest.mark.integration`` so you can
exclude them from CI with ``pytest -m "not integration"``.
"""

from __future__ import annotations

import base64
import os
import sys
from pathlib import Path
from typing import Any, Dict, Generator, List

import pytest

# ---------------------------------------------------------------------------
# Import SDK
# ---------------------------------------------------------------------------
from docling_serve_sdk import (
    ChunkDocumentResponse,
    ChunkedDocumentResultItem,
    ClearResponse,
    ConvertDocumentResponse,
    ConvertDocumentsRequest,
    ConvertDocumentsRequestOptions,
    DoclingClient,
    ExportResult,
    HealthCheckResponse,
    HierarchicalChunkerOptions,
    HierarchicalChunkerOptionsDocumentsRequest,
    HybridChunkerOptions,
    HybridChunkerOptionsDocumentsRequest,
    HttpSourceRequest,
    FileSourceRequest,
    OutputFormat,
    TaskStatusResponse,
)

# ---------------------------------------------------------------------------
# Configuration from environment
# ---------------------------------------------------------------------------
SERVER_URL = os.environ.get("DOCLING_SERVER_URL", "http://localhost:5001")
API_KEY = os.environ.get("DOCLING_API_KEY", None)

# Path to a small test PDF — create the fixtures directory and drop a PDF there
FIXTURES_DIR = Path(__file__).parent / "fixtures"
SAMPLE_PDF = FIXTURES_DIR / "sample.pdf"

# A small publicly-accessible PDF for URL-based tests
# (arXiv abstract page renders as a simple PDF)
SAMPLE_URL = "https://arxiv.org/pdf/2408.09869"

# Mark every test in this module as integration
pytestmark = pytest.mark.integration


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _require_sample_pdf() -> Path:
    """Skip the test if no sample PDF is available."""
    if not SAMPLE_PDF.exists():
        pytest.skip(
            f"No test PDF at {SAMPLE_PDF}. "
            f"Create tests/fixtures/sample.pdf to enable file-based tests."
        )
    return SAMPLE_PDF


def _make_file_source(pdf_path: Path) -> FileSourceRequest:
    """Create a FileSourceRequest from a local PDF."""
    raw = pdf_path.read_bytes()
    return FileSourceRequest(
        base64_string=base64.b64encode(raw).decode(),
        filename=pdf_path.name,
    )


def _log_ok(name: str, detail: str = "") -> None:
    extra = f" — {detail}" if detail else ""
    print(f"  ✅ {name}{extra}")


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

@pytest.fixture(scope="module")
def client() -> DoclingClient:
    """Create a client for the test session."""
    c = DoclingClient(
        SERVER_URL,
        api_key=API_KEY,
        timeout=120.0,
        max_retries=2,
    )
    return c


@pytest.fixture(scope="module")
def cm_client() -> Generator[DoclingClient, None, None]:
    """Client used via context manager (connection pooling)."""
    c = DoclingClient(SERVER_URL, api_key=API_KEY, timeout=120.0)
    with c:
        yield c


# ===================================================================
# 1. HEALTH & INFO
# ===================================================================

class TestHealthAndInfo:
    """Test GET /health and GET /version."""

    def test_health_check(self, client: DoclingClient) -> None:
        resp = client.health_check()
        assert isinstance(resp, HealthCheckResponse)
        assert resp.status == "ok"
        _log_ok("health_check", f"status={resp.status}")

    @pytest.mark.asyncio
    async def test_health_check_async(self, client: DoclingClient) -> None:
        async with DoclingClient(SERVER_URL, api_key=API_KEY, timeout=60) as c:
            resp = await c.health_check_async()
        assert isinstance(resp, HealthCheckResponse)
        assert resp.status == "ok"
        _log_ok("health_check_async", f"status={resp.status}")

    def test_version_info(self, client: DoclingClient) -> None:
        resp = client.version_info()
        assert isinstance(resp, dict)
        # Server should return at least some version info
        _log_ok("version_info", f"keys={list(resp.keys())}")

    @pytest.mark.asyncio
    async def test_version_info_async(self, client: DoclingClient) -> None:
        async with DoclingClient(SERVER_URL, api_key=API_KEY, timeout=60) as c:
            resp = await c.version_info_async()
        assert isinstance(resp, dict)
        _log_ok("version_info_async", f"keys={list(resp.keys())}")


# ===================================================================
# 2. DOCUMENT CONVERSION — Source (URL)
# ===================================================================

class TestConvertSource:
    """Test POST /v1/convert/source with HTTP URL source."""

    def test_convert_source_url_to_md(self, client: DoclingClient) -> None:
        """Convert a public PDF URL to Markdown."""
        req = ConvertDocumentsRequest(
            sources=[HttpSourceRequest(url=SAMPLE_URL)],
            options=ConvertDocumentsRequestOptions(
                to_formats=[OutputFormat.MD],
            ),
        )
        resp = client.convert_source(req)
        assert isinstance(resp, ConvertDocumentResponse)
        assert resp.document.md_content is not None
        assert len(resp.document.md_content) > 0
        assert resp.processing_time > 0
        _log_ok("convert_source (URL→MD)", f"{len(resp.document.md_content)} chars, {resp.processing_time:.2f}s")

    def test_convert_source_url_to_json(self, client: DoclingClient) -> None:
        """Convert a public PDF URL to JSON."""
        req = ConvertDocumentsRequest(
            sources=[HttpSourceRequest(url=SAMPLE_URL)],
            options=ConvertDocumentsRequestOptions(
                to_formats=[OutputFormat.JSON],
            ),
        )
        resp = client.convert_source(req)
        assert isinstance(resp, ConvertDocumentResponse)
        assert resp.document.json_content is not None
        _log_ok("convert_source (URL→JSON)", f"json_content keys={list(resp.document.json_content.keys())[:5]}")

    def test_convert_source_url_to_html(self, client: DoclingClient) -> None:
        """Convert a public PDF URL to HTML."""
        req = ConvertDocumentsRequest(
            sources=[HttpSourceRequest(url=SAMPLE_URL)],
            options=ConvertDocumentsRequestOptions(
                to_formats=[OutputFormat.HTML],
            ),
        )
        resp = client.convert_source(req)
        assert isinstance(resp, ConvertDocumentResponse)
        assert resp.document.html_content is not None
        assert "<" in resp.document.html_content
        _log_ok("convert_source (URL→HTML)", f"{len(resp.document.html_content)} chars")

    def test_convert_source_url_to_text(self, client: DoclingClient) -> None:
        """Convert a public PDF URL to plain text."""
        req = ConvertDocumentsRequest(
            sources=[HttpSourceRequest(url=SAMPLE_URL)],
            options=ConvertDocumentsRequestOptions(
                to_formats=[OutputFormat.TEXT],
            ),
        )
        resp = client.convert_source(req)
        assert isinstance(resp, ConvertDocumentResponse)
        assert resp.document.text_content is not None
        _log_ok("convert_source (URL→TEXT)", f"{len(resp.document.text_content)} chars")

    @pytest.mark.asyncio
    async def test_convert_source_async(self, client: DoclingClient) -> None:
        """Async variant of source conversion."""
        req = ConvertDocumentsRequest(
            sources=[HttpSourceRequest(url=SAMPLE_URL)],
            options=ConvertDocumentsRequestOptions(to_formats=[OutputFormat.MD]),
        )
        async with DoclingClient(SERVER_URL, api_key=API_KEY, timeout=120) as c:
            resp = await c.convert_source_async(req)
        assert isinstance(resp, ConvertDocumentResponse)
        assert resp.document.md_content is not None
        _log_ok("convert_source_async", f"{len(resp.document.md_content)} chars")


# ===================================================================
# 3. DOCUMENT CONVERSION — Source (Base64 file)
# ===================================================================

class TestConvertSourceBase64:
    """Test POST /v1/convert/source with base64-encoded file."""

    def test_convert_base64_pdf(self, client: DoclingClient) -> None:
        pdf_path = _require_sample_pdf()
        req = ConvertDocumentsRequest(
            sources=[_make_file_source(pdf_path)],
            options=ConvertDocumentsRequestOptions(to_formats=[OutputFormat.MD]),
        )
        resp = client.convert_source(req)
        assert isinstance(resp, ConvertDocumentResponse)
        assert resp.document.md_content is not None
        assert len(resp.document.md_content) > 0
        _log_ok("convert_source (base64→MD)", f"{len(resp.document.md_content)} chars")


# ===================================================================
# 4. DOCUMENT CONVERSION — File upload
# ===================================================================

class TestConvertFile:
    """Test POST /v1/convert/file with multipart upload."""

    def test_convert_file_to_md(self, client: DoclingClient) -> None:
        pdf_path = _require_sample_pdf()
        resp = client.convert_file(pdf_path)
        assert isinstance(resp, ConvertDocumentResponse)
        assert resp.document.md_content is not None
        assert resp.processing_time > 0
        _log_ok("convert_file (→MD)", f"{len(resp.document.md_content)} chars, {resp.processing_time:.2f}s")

    def test_convert_file_to_json(self, client: DoclingClient) -> None:
        pdf_path = _require_sample_pdf()
        resp = client.convert_file(
            pdf_path,
            options=ConvertDocumentsRequestOptions(to_formats=[OutputFormat.JSON]),
        )
        assert isinstance(resp, ConvertDocumentResponse)
        assert resp.document.json_content is not None
        _log_ok("convert_file (→JSON)", f"keys={list(resp.document.json_content.keys())[:5]}")

    def test_convert_file_with_ocr_disabled(self, client: DoclingClient) -> None:
        pdf_path = _require_sample_pdf()
        resp = client.convert_file(
            pdf_path,
            options=ConvertDocumentsRequestOptions(do_ocr=False),
        )
        assert isinstance(resp, ConvertDocumentResponse)
        _log_ok("convert_file (no OCR)", f"processing_time={resp.processing_time:.2f}s")

    def test_convert_file_binary_io(self, client: DoclingClient) -> None:
        """Test passing an open file handle instead of a path."""
        pdf_path = _require_sample_pdf()
        with open(pdf_path, "rb") as fh:
            resp = client.convert_file(fh)
        assert isinstance(resp, ConvertDocumentResponse)
        assert resp.document.md_content is not None
        _log_ok("convert_file (BinaryIO)", f"{len(resp.document.md_content)} chars")

    @pytest.mark.asyncio
    async def test_convert_file_async(self, client: DoclingClient) -> None:
        pdf_path = _require_sample_pdf()
        async with DoclingClient(SERVER_URL, api_key=API_KEY, timeout=120) as c:
            resp = await c.convert_file_async(pdf_path)
        assert isinstance(resp, ConvertDocumentResponse)
        assert resp.document.md_content is not None
        _log_ok("convert_file_async", f"{len(resp.document.md_content)} chars")

    def test_convert_file_multiple_formats(self, client: DoclingClient) -> None:
        """Request both MD and JSON output at once."""
        pdf_path = _require_sample_pdf()
        resp = client.convert_file(
            pdf_path,
            options=ConvertDocumentsRequestOptions(
                to_formats=[OutputFormat.MD, OutputFormat.JSON],
            ),
        )
        assert isinstance(resp, ConvertDocumentResponse)
        assert resp.document.md_content is not None
        assert resp.document.json_content is not None
        _log_ok("convert_file (MD+JSON)", "both formats returned")


# ===================================================================
# 5. ASYNC TASK — Source
# ===================================================================

class TestAsyncTaskSource:
    """Test POST /v1/convert/source/async + polling + result."""

    def test_async_source_full_flow(self, client: DoclingClient) -> None:
        """Submit → poll → wait → result."""
        req = ConvertDocumentsRequest(
            sources=[HttpSourceRequest(url=SAMPLE_URL)],
            options=ConvertDocumentsRequestOptions(to_formats=[OutputFormat.MD]),
        )
        # Submit
        task = client.convert_source_async_task(req)
        assert isinstance(task, TaskStatusResponse)
        assert task.task_id
        _log_ok("convert_source_async_task", f"task_id={task.task_id}")

        # Poll once
        status = client.task_status_poll(task.task_id)
        assert isinstance(status, TaskStatusResponse)
        _log_ok("task_status_poll", f"status={status.task_status}")

        # Wait for completion
        result = client.wait_for_task(task.task_id, poll_interval=2.0, timeout=120.0)
        assert isinstance(result, dict)
        _log_ok("wait_for_task", f"result keys={list(result.keys())[:5]}")

    @pytest.mark.asyncio
    async def test_async_source_full_flow_async(self, client: DoclingClient) -> None:
        req = ConvertDocumentsRequest(
            sources=[HttpSourceRequest(url=SAMPLE_URL)],
            options=ConvertDocumentsRequestOptions(to_formats=[OutputFormat.MD]),
        )
        async with DoclingClient(SERVER_URL, api_key=API_KEY, timeout=120) as c:
            task = await c.convert_source_async_task_async(req)
            assert task.task_id
            result = await c.wait_for_task_async(task.task_id, poll_interval=2.0, timeout=120.0)
            assert isinstance(result, dict)
        _log_ok("async_source_full_flow_async", f"task_id={task.task_id}")


# ===================================================================
# 6. ASYNC TASK — File
# ===================================================================

class TestAsyncTaskFile:
    """Test POST /v1/convert/file/async + polling + result."""

    def test_async_file_full_flow(self, client: DoclingClient) -> None:
        pdf_path = _require_sample_pdf()
        task = client.convert_file_async_task(pdf_path)
        assert isinstance(task, TaskStatusResponse)
        assert task.task_id
        _log_ok("convert_file_async_task", f"task_id={task.task_id}")

        result = client.wait_for_task(task.task_id, poll_interval=2.0, timeout=120.0)
        assert isinstance(result, dict)
        _log_ok("wait_for_task (file)", f"result keys={list(result.keys())[:5]}")

    @pytest.mark.asyncio
    async def test_async_file_full_flow_async(self, client: DoclingClient) -> None:
        pdf_path = _require_sample_pdf()
        async with DoclingClient(SERVER_URL, api_key=API_KEY, timeout=120) as c:
            task = await c.convert_file_async_task_async(pdf_path)
            assert task.task_id
            result = await c.wait_for_task_async(task.task_id, poll_interval=2.0, timeout=120.0)
            assert isinstance(result, dict)
        _log_ok("async_file_full_flow_async", f"task_id={task.task_id}")


# ===================================================================
# 7. CHUNKING — Hybrid (source)
# ===================================================================

class TestChunkHybridSource:
    """Test POST /v1/chunk/hybrid/source."""

    def test_chunk_hybrid_source_default(self, client: DoclingClient) -> None:
        req = HybridChunkerOptionsDocumentsRequest(
            sources=[HttpSourceRequest(url=SAMPLE_URL)],
        )
        resp = client.chunk_hybrid_source(req)
        assert isinstance(resp, ChunkDocumentResponse)
        assert len(resp.chunks) > 0
        assert resp.processing_time > 0
        chunk = resp.chunks[0]
        assert isinstance(chunk, ChunkedDocumentResultItem)
        assert chunk.text
        _log_ok("chunk_hybrid_source", f"{len(resp.chunks)} chunks, {resp.processing_time:.2f}s")

    def test_chunk_hybrid_source_with_options(self, client: DoclingClient) -> None:
        req = HybridChunkerOptionsDocumentsRequest(
            sources=[HttpSourceRequest(url=SAMPLE_URL)],
            chunking_options=HybridChunkerOptions(
                max_tokens=128,
                include_raw_text=True,
            ),
        )
        resp = client.chunk_hybrid_source(req)
        assert isinstance(resp, ChunkDocumentResponse)
        assert len(resp.chunks) > 0
        # With include_raw_text=True, chunks should have raw_text
        _log_ok("chunk_hybrid_source (opts)", f"{len(resp.chunks)} chunks")

    def test_chunk_hybrid_source_include_doc(self, client: DoclingClient) -> None:
        req = HybridChunkerOptionsDocumentsRequest(
            sources=[HttpSourceRequest(url=SAMPLE_URL)],
            include_converted_doc=True,
        )
        resp = client.chunk_hybrid_source(req)
        assert isinstance(resp, ChunkDocumentResponse)
        assert len(resp.chunks) > 0
        # With include_converted_doc=True, documents list should be populated
        if len(resp.documents) > 0:
            assert isinstance(resp.documents[0], ExportResult)
            _log_ok("chunk_hybrid_source (include_doc)", f"{len(resp.documents)} docs returned")
        else:
            _log_ok("chunk_hybrid_source (include_doc)", "no docs (server may omit)")

    @pytest.mark.asyncio
    async def test_chunk_hybrid_source_async(self, client: DoclingClient) -> None:
        req = HybridChunkerOptionsDocumentsRequest(
            sources=[HttpSourceRequest(url=SAMPLE_URL)],
        )
        async with DoclingClient(SERVER_URL, api_key=API_KEY, timeout=120) as c:
            resp = await c.chunk_hybrid_source_async(req)
        assert isinstance(resp, ChunkDocumentResponse)
        assert len(resp.chunks) > 0
        _log_ok("chunk_hybrid_source_async", f"{len(resp.chunks)} chunks")


# ===================================================================
# 8. CHUNKING — Hybrid (file)
# ===================================================================

class TestChunkHybridFile:
    """Test POST /v1/chunk/hybrid/file."""

    def test_chunk_hybrid_file_default(self, client: DoclingClient) -> None:
        pdf_path = _require_sample_pdf()
        resp = client.chunk_hybrid_file(pdf_path)
        assert isinstance(resp, ChunkDocumentResponse)
        assert len(resp.chunks) > 0
        assert resp.processing_time > 0
        _log_ok("chunk_hybrid_file", f"{len(resp.chunks)} chunks, {resp.processing_time:.2f}s")

    def test_chunk_hybrid_file_with_options(self, client: DoclingClient) -> None:
        pdf_path = _require_sample_pdf()
        resp = client.chunk_hybrid_file(
            pdf_path,
            chunking_options=HybridChunkerOptions(max_tokens=64),
        )
        assert isinstance(resp, ChunkDocumentResponse)
        assert len(resp.chunks) > 0
        _log_ok("chunk_hybrid_file (max_tokens=64)", f"{len(resp.chunks)} chunks")

    @pytest.mark.asyncio
    async def test_chunk_hybrid_file_async(self, client: DoclingClient) -> None:
        pdf_path = _require_sample_pdf()
        async with DoclingClient(SERVER_URL, api_key=API_KEY, timeout=120) as c:
            resp = await c.chunk_hybrid_file_async(pdf_path)
        assert isinstance(resp, ChunkDocumentResponse)
        assert len(resp.chunks) > 0
        _log_ok("chunk_hybrid_file_async", f"{len(resp.chunks)} chunks")


# ===================================================================
# 9. CHUNKING — Hybrid (async task)
# ===================================================================

class TestChunkHybridAsyncTask:
    """Test POST /v1/chunk/hybrid/source/async and file/async."""

    def test_chunk_hybrid_source_async_task(self, client: DoclingClient) -> None:
        req = HybridChunkerOptionsDocumentsRequest(
            sources=[HttpSourceRequest(url=SAMPLE_URL)],
        )
        task = client.chunk_hybrid_source_async_task(req)
        assert isinstance(task, TaskStatusResponse)
        assert task.task_id
        result = client.wait_for_task(task.task_id, poll_interval=2.0, timeout=120.0)
        assert isinstance(result, dict)
        _log_ok("chunk_hybrid_source_async_task", f"task_id={task.task_id}")

    def test_chunk_hybrid_file_async_task(self, client: DoclingClient) -> None:
        pdf_path = _require_sample_pdf()
        task = client.chunk_hybrid_file_async_task(pdf_path)
        assert isinstance(task, TaskStatusResponse)
        assert task.task_id
        result = client.wait_for_task(task.task_id, poll_interval=2.0, timeout=120.0)
        assert isinstance(result, dict)
        _log_ok("chunk_hybrid_file_async_task", f"task_id={task.task_id}")

    @pytest.mark.asyncio
    async def test_chunk_hybrid_source_async_task_async(self, client: DoclingClient) -> None:
        req = HybridChunkerOptionsDocumentsRequest(
            sources=[HttpSourceRequest(url=SAMPLE_URL)],
        )
        async with DoclingClient(SERVER_URL, api_key=API_KEY, timeout=120) as c:
            task = await c.chunk_hybrid_source_async_task_async(req)
            assert task.task_id
            result = await c.wait_for_task_async(task.task_id, poll_interval=2.0, timeout=120.0)
            assert isinstance(result, dict)
        _log_ok("chunk_hybrid_source_async_task_async", f"task_id={task.task_id}")


# ===================================================================
# 10. CHUNKING — Hierarchical (source)
# ===================================================================

class TestChunkHierarchicalSource:
    """Test POST /v1/chunk/hierarchical/source."""

    def test_chunk_hierarchical_source_default(self, client: DoclingClient) -> None:
        req = HierarchicalChunkerOptionsDocumentsRequest(
            sources=[HttpSourceRequest(url=SAMPLE_URL)],
        )
        resp = client.chunk_hierarchical_source(req)
        assert isinstance(resp, ChunkDocumentResponse)
        assert len(resp.chunks) > 0
        assert resp.processing_time > 0
        _log_ok("chunk_hierarchical_source", f"{len(resp.chunks)} chunks, {resp.processing_time:.2f}s")

    def test_chunk_hierarchical_source_with_options(self, client: DoclingClient) -> None:
        req = HierarchicalChunkerOptionsDocumentsRequest(
            sources=[HttpSourceRequest(url=SAMPLE_URL)],
            chunking_options=HierarchicalChunkerOptions(
                include_raw_text=True,
                use_markdown_tables=True,
            ),
        )
        resp = client.chunk_hierarchical_source(req)
        assert isinstance(resp, ChunkDocumentResponse)
        assert len(resp.chunks) > 0
        _log_ok("chunk_hierarchical_source (opts)", f"{len(resp.chunks)} chunks")

    @pytest.mark.asyncio
    async def test_chunk_hierarchical_source_async(self, client: DoclingClient) -> None:
        req = HierarchicalChunkerOptionsDocumentsRequest(
            sources=[HttpSourceRequest(url=SAMPLE_URL)],
        )
        async with DoclingClient(SERVER_URL, api_key=API_KEY, timeout=120) as c:
            resp = await c.chunk_hierarchical_source_async(req)
        assert isinstance(resp, ChunkDocumentResponse)
        assert len(resp.chunks) > 0
        _log_ok("chunk_hierarchical_source_async", f"{len(resp.chunks)} chunks")


# ===================================================================
# 11. CHUNKING — Hierarchical (file)
# ===================================================================

class TestChunkHierarchicalFile:
    """Test POST /v1/chunk/hierarchical/file."""

    def test_chunk_hierarchical_file_default(self, client: DoclingClient) -> None:
        pdf_path = _require_sample_pdf()
        resp = client.chunk_hierarchical_file(pdf_path)
        assert isinstance(resp, ChunkDocumentResponse)
        assert len(resp.chunks) > 0
        _log_ok("chunk_hierarchical_file", f"{len(resp.chunks)} chunks, {resp.processing_time:.2f}s")

    @pytest.mark.asyncio
    async def test_chunk_hierarchical_file_async(self, client: DoclingClient) -> None:
        pdf_path = _require_sample_pdf()
        async with DoclingClient(SERVER_URL, api_key=API_KEY, timeout=120) as c:
            resp = await c.chunk_hierarchical_file_async(pdf_path)
        assert isinstance(resp, ChunkDocumentResponse)
        assert len(resp.chunks) > 0
        _log_ok("chunk_hierarchical_file_async", f"{len(resp.chunks)} chunks")


# ===================================================================
# 12. CHUNKING — Hierarchical (async task)
# ===================================================================

class TestChunkHierarchicalAsyncTask:
    """Test POST /v1/chunk/hierarchical/source/async and file/async."""

    def test_chunk_hierarchical_source_async_task(self, client: DoclingClient) -> None:
        req = HierarchicalChunkerOptionsDocumentsRequest(
            sources=[HttpSourceRequest(url=SAMPLE_URL)],
        )
        task = client.chunk_hierarchical_source_async_task(req)
        assert isinstance(task, TaskStatusResponse)
        assert task.task_id
        result = client.wait_for_task(task.task_id, poll_interval=2.0, timeout=120.0)
        assert isinstance(result, dict)
        _log_ok("chunk_hierarchical_source_async_task", f"task_id={task.task_id}")

    def test_chunk_hierarchical_file_async_task(self, client: DoclingClient) -> None:
        pdf_path = _require_sample_pdf()
        task = client.chunk_hierarchical_file_async_task(pdf_path)
        assert isinstance(task, TaskStatusResponse)
        assert task.task_id
        result = client.wait_for_task(task.task_id, poll_interval=2.0, timeout=120.0)
        assert isinstance(result, dict)
        _log_ok("chunk_hierarchical_file_async_task", f"task_id={task.task_id}")

    @pytest.mark.asyncio
    async def test_chunk_hierarchical_file_async_task_async(self, client: DoclingClient) -> None:
        pdf_path = _require_sample_pdf()
        async with DoclingClient(SERVER_URL, api_key=API_KEY, timeout=120) as c:
            task = await c.chunk_hierarchical_file_async_task_async(pdf_path)
            assert task.task_id
            result = await c.wait_for_task_async(task.task_id, poll_interval=2.0, timeout=120.0)
            assert isinstance(result, dict)
        _log_ok("chunk_hierarchical_file_async_task_async", f"task_id={task.task_id}")


# ===================================================================
# 13. MANAGEMENT — Clear & Memory
# ===================================================================

class TestManagement:
    """Test clear and memory endpoints."""

    def test_clear_converters(self, client: DoclingClient) -> None:
        resp = client.clear_converters()
        assert isinstance(resp, ClearResponse)
        _log_ok("clear_converters", f"status={resp.status}")

    @pytest.mark.asyncio
    async def test_clear_converters_async(self, client: DoclingClient) -> None:
        async with DoclingClient(SERVER_URL, api_key=API_KEY, timeout=60) as c:
            resp = await c.clear_converters_async()
        assert isinstance(resp, ClearResponse)
        _log_ok("clear_converters_async", f"status={resp.status}")

    def test_clear_results(self, client: DoclingClient) -> None:
        resp = client.clear_results(older_than=0)
        assert isinstance(resp, ClearResponse)
        _log_ok("clear_results", f"status={resp.status}")

    @pytest.mark.asyncio
    async def test_clear_results_async(self, client: DoclingClient) -> None:
        async with DoclingClient(SERVER_URL, api_key=API_KEY, timeout=60) as c:
            resp = await c.clear_results_async(older_than=0)
        assert isinstance(resp, ClearResponse)
        _log_ok("clear_results_async", f"status={resp.status}")

    def test_memory_stats(self, client: DoclingClient) -> None:
        from docling_serve_sdk import DoclingAPIError
        try:
            resp = client.memory_stats()
            assert isinstance(resp, dict)
            _log_ok("memory_stats", f"keys={list(resp.keys())[:5]}")
        except DoclingAPIError as e:
            if e.status_code == 403:
                pytest.skip("Memory endpoints disabled on this server")
            raise

    @pytest.mark.asyncio
    async def test_memory_stats_async(self, client: DoclingClient) -> None:
        from docling_serve_sdk import DoclingAPIError
        try:
            async with DoclingClient(SERVER_URL, api_key=API_KEY, timeout=60) as c:
                resp = await c.memory_stats_async()
            assert isinstance(resp, dict)
            _log_ok("memory_stats_async", f"keys={list(resp.keys())[:5]}")
        except DoclingAPIError as e:
            if e.status_code == 403:
                pytest.skip("Memory endpoints disabled on this server")
            raise

    def test_memory_counts(self, client: DoclingClient) -> None:
        from docling_serve_sdk import DoclingAPIError
        try:
            resp = client.memory_counts()
            assert isinstance(resp, dict)
            _log_ok("memory_counts", f"keys={list(resp.keys())[:5]}")
        except DoclingAPIError as e:
            if e.status_code == 403:
                pytest.skip("Memory endpoints disabled on this server")
            raise

    @pytest.mark.asyncio
    async def test_memory_counts_async(self, client: DoclingClient) -> None:
        from docling_serve_sdk import DoclingAPIError
        try:
            async with DoclingClient(SERVER_URL, api_key=API_KEY, timeout=60) as c:
                resp = await c.memory_counts_async()
            assert isinstance(resp, dict)
            _log_ok("memory_counts_async", f"keys={list(resp.keys())[:5]}")
        except DoclingAPIError as e:
            if e.status_code == 403:
                pytest.skip("Memory endpoints disabled on this server")
            raise


# ===================================================================
# 14. ENTERPRISE FEATURES — Hooks, Context Manager, Properties
# ===================================================================

class TestEnterpriseFeatures:
    """Test enterprise features: hooks, connection pooling, properties."""

    def test_context_manager_sync(self, client: DoclingClient) -> None:
        """Verify sync context manager works for connection pooling."""
        with DoclingClient(SERVER_URL, api_key=API_KEY, timeout=60) as c:
            resp = c.health_check()
            assert resp.status == "ok"
            # Second call reuses the pooled connection
            resp2 = c.health_check()
            assert resp2.status == "ok"
        _log_ok("context_manager (sync)", "pooled 2 requests")

    @pytest.mark.asyncio
    async def test_context_manager_async(self, client: DoclingClient) -> None:
        """Verify async context manager works for connection pooling."""
        async with DoclingClient(SERVER_URL, api_key=API_KEY, timeout=60) as c:
            resp = await c.health_check_async()
            assert resp.status == "ok"
            resp2 = await c.health_check_async()
            assert resp2.status == "ok"
        _log_ok("context_manager (async)", "pooled 2 requests")

    def test_on_request_hook_fires(self, client: DoclingClient) -> None:
        """Verify on_request hook is called."""
        hook_calls: List[str] = []

        def on_req(rid: str, method: str, path: str, body: Any) -> None:
            hook_calls.append(f"{method} {path}")

        c = DoclingClient(SERVER_URL, api_key=API_KEY, timeout=60, on_request=on_req)
        c.health_check()
        assert len(hook_calls) == 1
        assert "GET" in hook_calls[0] and "/health" in hook_calls[0]
        _log_ok("on_request hook", f"captured: {hook_calls[0]}")

    def test_on_response_hook_fires(self, client: DoclingClient) -> None:
        """Verify on_response hook captures status code and duration."""
        hook_calls: List[Dict[str, Any]] = []

        def on_resp(rid: str, method: str, path: str, status: int, dur: float) -> None:
            hook_calls.append({"status": status, "duration": dur, "path": path})

        c = DoclingClient(SERVER_URL, api_key=API_KEY, timeout=60, on_response=on_resp)
        c.health_check()
        assert len(hook_calls) == 1
        assert hook_calls[0]["status"] == 200
        assert hook_calls[0]["duration"] > 0
        _log_ok("on_response hook", f"status={hook_calls[0]['status']}, dur={hook_calls[0]['duration']:.4f}s")

    def test_client_properties(self) -> None:
        """Verify public read-only properties."""
        c = DoclingClient(
            SERVER_URL,
            api_key="test-key",
            timeout=42.0,
            max_retries=5,
            verify_ssl=False,
        )
        assert c.base_url == SERVER_URL.rstrip("/")
        assert c.timeout == 42.0
        assert c.max_retries == 5
        assert c.verify_ssl is False
        headers = c.headers
        assert "X-Api-Key" in headers
        assert headers["X-Api-Key"] == "test-key"
        # headers is a copy — mutating it shouldn't affect the client
        headers["X-Api-Key"] = "hacked"
        assert c.headers["X-Api-Key"] == "test-key"
        _log_ok("client_properties", "all props correct, headers is immutable copy")

    def test_repr_redacts_api_key(self) -> None:
        c = DoclingClient(SERVER_URL, api_key="super-secret-key")
        r = repr(c)
        assert "super-secret-key" not in r
        assert "***" in r
        _log_ok("repr", "API key redacted")

    def test_explicit_close(self) -> None:
        """Verify explicit close() doesn't crash."""
        c = DoclingClient(SERVER_URL, api_key=API_KEY, timeout=60)
        with c:
            c.health_check()
        # close() after context manager exit should be safe
        c.close()
        _log_ok("explicit_close", "no error on double close")


# ===================================================================
# 15. RESPONSE MODEL VALIDATION
# ===================================================================

class TestResponseModelValidation:
    """Verify server responses deserialize into correct Pydantic models."""

    def test_convert_response_fields(self, client: DoclingClient) -> None:
        """ConvertDocumentResponse has typed errors + timings."""
        pdf_path = _require_sample_pdf()
        # Late in the suite the server may be under load; tolerate transient failures.
        heavy_client = DoclingClient(
            SERVER_URL, api_key=API_KEY, timeout=300.0, max_retries=3,
        )
        try:
            raw_resp = heavy_client.convert_file(pdf_path)
        except Exception as exc:
            pytest.skip(f"Server unavailable at end of suite: {exc}")
        assert isinstance(raw_resp, ConvertDocumentResponse)
        resp = raw_resp
        assert isinstance(resp.errors, list)
        assert isinstance(resp.timings, dict)
        _log_ok("response fields", f"errors={len(resp.errors)}, timings_keys={list(resp.timings.keys())[:3]}")

    def test_chunk_response_items_are_typed(self, client: DoclingClient) -> None:
        """ChunkDocumentResponse.chunks should be List[ChunkedDocumentResultItem]."""
        pdf_path = _require_sample_pdf()
        heavy_client = DoclingClient(
            SERVER_URL, api_key=API_KEY, timeout=300.0, max_retries=3,
        )
        try:
            resp = heavy_client.chunk_hybrid_file(pdf_path)
        except Exception as exc:
            pytest.skip(f"Server unavailable at end of suite: {exc}")
        assert len(resp.chunks) > 0
        chunk = resp.chunks[0]
        assert isinstance(chunk, ChunkedDocumentResultItem)
        assert isinstance(chunk.filename, str)
        assert isinstance(chunk.chunk_index, int)
        assert isinstance(chunk.text, str)
        assert isinstance(chunk.doc_items, list)
        _log_ok("chunk items typed", f"filename={chunk.filename}, idx={chunk.chunk_index}")


# ===================================================================
# 16. ERROR HANDLING
# ===================================================================

class TestErrorHandling:
    """Test that SDK raises proper typed exceptions on errors."""

    def test_invalid_url_source_returns_error(self, client: DoclingClient) -> None:
        """Server should reject an invalid URL."""
        from docling_serve_sdk import DoclingAPIError
        req = ConvertDocumentsRequest(
            sources=[HttpSourceRequest(url="https://this-domain-does-not-exist-xyz.invalid/doc.pdf")],
        )
        with pytest.raises((DoclingAPIError, Exception)):
            client.convert_source(req)
        _log_ok("error: invalid URL", "raised exception as expected")

    def test_connection_error_on_bad_server(self) -> None:
        """SDK should raise DoclingConnectionError for unreachable server."""
        from docling_serve_sdk import DoclingConnectionError, DoclingTimeoutError
        c = DoclingClient("http://127.0.0.1:59999", timeout=3.0, max_retries=0)
        with pytest.raises((DoclingConnectionError, DoclingTimeoutError, Exception)):
            c.health_check()
        _log_ok("error: connection refused", "raised exception as expected")


# ===================================================================
# MAIN — Run as script for quick smoke test
# ===================================================================

if __name__ == "__main__":
    print(f"\n{'='*60}")
    print(f"  Docling Serve SDK — Integration Test Suite")
    print(f"  Server: {SERVER_URL}")
    print(f"  API Key: {'***' if API_KEY else 'None'}")
    print(f"  Sample PDF: {SAMPLE_PDF} ({'exists' if SAMPLE_PDF.exists() else 'MISSING'})")
    print(f"{'='*60}\n")
    sys.exit(pytest.main([__file__, "-v", "--tb=short", "-x"]))
