"""Unit tests for vision structured output extraction."""
from __future__ import annotations

import base64
import json
import os
import pathlib
from unittest.mock import MagicMock, patch

import pytest

# ---------------------------------------------------------------------------
# Load Azure env vars from .env if available (for integration tests)
# ---------------------------------------------------------------------------
_PROJECT_DIR = pathlib.Path(__file__).resolve().parent.parent
_ENV_FILE = _PROJECT_DIR / ".env"

# Test image from grab-receipt-manager (used in integration tests only)
_TEST_IMAGE_DIR = pathlib.Path(
    "/Users/dewa/Documents/srx/codes/HOME/z-legacy-grab-receipt-manager/test-data"
)


def _load_env():
    if not _ENV_FILE.exists():
        return
    for line in _ENV_FILE.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        if "=" not in line:
            continue
        key, _, value = line.partition("=")
        key = key.strip()
        value = value.strip()
        if key.startswith("AZURE_OPENAI_") or key == "OPENAI_API_KEY":
            os.environ.setdefault(key, value)


_load_env()

AZURE_AVAILABLE = bool(os.getenv("AZURE_OPENAI_ENDPOINT"))
skip_no_azure = pytest.mark.skipif(
    not AZURE_AVAILABLE, reason="Azure OpenAI env vars not set"
)

# ---------------------------------------------------------------------------
# Shared fixtures
# ---------------------------------------------------------------------------

SAMPLE_SCHEMA = {
    "type": "object",
    "properties": {
        "item_name": {"type": "string", "description": "Name of the item"},
        "amount": {"type": "number", "description": "Total amount"},
        "currency": {"type": "string", "description": "Currency code"},
    },
    "required": ["item_name", "amount"],
}

SAMPLE_RESPONSE = {"item_name": "Coffee", "amount": 25000, "currency": "IDR"}

SAMPLE_IMAGE_URL = "https://example.com/receipt.jpg"
SAMPLE_BASE64 = "data:image/png;base64," + base64.b64encode(b"\x89PNG\r\n").decode()


def _mock_completion(content: dict) -> MagicMock:
    """Build a mock ChatCompletion response."""
    mock_resp = MagicMock()
    mock_resp.choices = [MagicMock()]
    mock_resp.choices[0].message.content = json.dumps(content)
    return mock_resp


# ---------------------------------------------------------------------------
# Unit tests (mocked – no real API calls)
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_extract_structured_vision_url():
    """Verify message structure with a URL image."""
    from srx_lib_llm.vision import extract_structured_vision

    mock_client = MagicMock()
    mock_client.chat.completions.create.return_value = _mock_completion(SAMPLE_RESPONSE)

    with (
        patch("srx_lib_llm.vision.is_azure_openai", return_value=False),
        patch("srx_lib_llm.vision.get_effective_model", return_value="gpt-4.1-mini"),
        patch("srx_lib_llm.vision.get_openai_client_kwargs", return_value={}),
        patch("openai.OpenAI", return_value=mock_client),
    ):
        result = await extract_structured_vision(
            images=[SAMPLE_IMAGE_URL],
            instruction="Extract the receipt data.",
            json_schema=SAMPLE_SCHEMA,
            schema_name="Receipt",
        )

    # Verify result
    assert result.item_name == "Coffee"
    assert result.amount == 25000

    # Verify the call was made with correct structure
    call_kwargs = mock_client.chat.completions.create.call_args.kwargs
    assert call_kwargs["model"] == "gpt-4.1-mini"
    assert call_kwargs["response_format"]["type"] == "json_schema"
    assert call_kwargs["response_format"]["json_schema"]["strict"] is True

    # Verify user message has both text and image_url parts
    user_msg = call_kwargs["messages"][1]
    assert user_msg["role"] == "user"
    assert isinstance(user_msg["content"], list)
    assert user_msg["content"][0]["type"] == "text"
    assert user_msg["content"][1]["type"] == "image_url"
    assert user_msg["content"][1]["image_url"]["url"] == SAMPLE_IMAGE_URL
    assert user_msg["content"][1]["image_url"]["detail"] == "auto"


@pytest.mark.asyncio
async def test_extract_structured_vision_base64():
    """Verify base64 data URIs are passed through without transformation."""
    from srx_lib_llm.vision import extract_structured_vision

    mock_client = MagicMock()
    mock_client.chat.completions.create.return_value = _mock_completion(SAMPLE_RESPONSE)

    with (
        patch("srx_lib_llm.vision.is_azure_openai", return_value=False),
        patch("srx_lib_llm.vision.get_effective_model", return_value="gpt-4.1-mini"),
        patch("srx_lib_llm.vision.get_openai_client_kwargs", return_value={}),
        patch("openai.OpenAI", return_value=mock_client),
    ):
        result = await extract_structured_vision(
            images=[SAMPLE_BASE64],
            instruction="Extract data.",
            json_schema=SAMPLE_SCHEMA,
        )

    call_kwargs = mock_client.chat.completions.create.call_args.kwargs
    img_part = call_kwargs["messages"][1]["content"][1]
    assert img_part["image_url"]["url"] == SAMPLE_BASE64


@pytest.mark.asyncio
async def test_extract_structured_vision_multiple_images():
    """Verify multiple images produce multiple image_url parts."""
    from srx_lib_llm.vision import extract_structured_vision

    mock_client = MagicMock()
    mock_client.chat.completions.create.return_value = _mock_completion(SAMPLE_RESPONSE)

    urls = ["https://example.com/page1.jpg", "https://example.com/page2.jpg"]

    with (
        patch("srx_lib_llm.vision.is_azure_openai", return_value=False),
        patch("srx_lib_llm.vision.get_effective_model", return_value="gpt-4.1-mini"),
        patch("srx_lib_llm.vision.get_openai_client_kwargs", return_value={}),
        patch("openai.OpenAI", return_value=mock_client),
    ):
        await extract_structured_vision(
            images=urls,
            instruction="Extract data from multi-page receipt.",
            json_schema=SAMPLE_SCHEMA,
        )

    call_kwargs = mock_client.chat.completions.create.call_args.kwargs
    content = call_kwargs["messages"][1]["content"]
    # 1 text part + 2 image_url parts
    assert len(content) == 3
    assert content[1]["image_url"]["url"] == urls[0]
    assert content[2]["image_url"]["url"] == urls[1]


@pytest.mark.asyncio
@pytest.mark.parametrize("detail", ["auto", "high", "low"])
async def test_extract_structured_vision_detail_levels(detail):
    """Verify the detail parameter is forwarded correctly."""
    from srx_lib_llm.vision import extract_structured_vision

    mock_client = MagicMock()
    mock_client.chat.completions.create.return_value = _mock_completion(SAMPLE_RESPONSE)

    with (
        patch("srx_lib_llm.vision.is_azure_openai", return_value=False),
        patch("srx_lib_llm.vision.get_effective_model", return_value="gpt-4.1-mini"),
        patch("srx_lib_llm.vision.get_openai_client_kwargs", return_value={}),
        patch("openai.OpenAI", return_value=mock_client),
    ):
        await extract_structured_vision(
            images=[SAMPLE_IMAGE_URL],
            instruction="Extract.",
            json_schema=SAMPLE_SCHEMA,
            detail=detail,
        )

    call_kwargs = mock_client.chat.completions.create.call_args.kwargs
    assert call_kwargs["messages"][1]["content"][1]["image_url"]["detail"] == detail


@pytest.mark.asyncio
async def test_extract_structured_vision_azure_client():
    """Verify Azure path uses create_openai_client."""
    from srx_lib_llm.vision import extract_structured_vision

    mock_client = MagicMock()
    mock_client.chat.completions.create.return_value = _mock_completion(SAMPLE_RESPONSE)

    with (
        patch("srx_lib_llm.vision.is_azure_openai", return_value=True),
        patch("srx_lib_llm.vision.get_effective_model", return_value="gpt-5-mini"),
        patch("srx_lib_llm.vision.create_openai_client", return_value=mock_client) as mock_factory,
    ):
        result = await extract_structured_vision(
            images=[SAMPLE_IMAGE_URL],
            instruction="Extract the receipt data.",
            json_schema=SAMPLE_SCHEMA,
        )

    mock_factory.assert_called_once_with(api_key=None, async_=False)
    assert result.item_name == "Coffee"


@pytest.mark.asyncio
async def test_extract_structured_vision_reasoning_effort():
    """Verify reasoning_effort is passed for GPT-5.x models."""
    from srx_lib_llm.vision import extract_structured_vision

    mock_client = MagicMock()
    mock_client.chat.completions.create.return_value = _mock_completion(SAMPLE_RESPONSE)

    with (
        patch("srx_lib_llm.vision.is_azure_openai", return_value=True),
        patch("srx_lib_llm.vision.get_effective_model", return_value="gpt-5-mini"),
        patch("srx_lib_llm.vision.create_openai_client", return_value=mock_client),
    ):
        await extract_structured_vision(
            images=[SAMPLE_IMAGE_URL],
            instruction="Extract.",
            json_schema=SAMPLE_SCHEMA,
            reasoning_effort="medium",
        )

    call_kwargs = mock_client.chat.completions.create.call_args.kwargs
    assert call_kwargs["reasoning_effort"] == "medium"


@pytest.mark.asyncio
async def test_extract_structured_vision_reasoning_none_not_passed():
    """Verify reasoning_effort='none' is NOT passed (API doesn't accept it)."""
    from srx_lib_llm.vision import extract_structured_vision

    mock_client = MagicMock()
    mock_client.chat.completions.create.return_value = _mock_completion(SAMPLE_RESPONSE)

    with (
        patch("srx_lib_llm.vision.is_azure_openai", return_value=True),
        patch("srx_lib_llm.vision.get_effective_model", return_value="gpt-5-mini"),
        patch("srx_lib_llm.vision.create_openai_client", return_value=mock_client),
    ):
        await extract_structured_vision(
            images=[SAMPLE_IMAGE_URL],
            instruction="Extract.",
            json_schema=SAMPLE_SCHEMA,
            reasoning_effort="none",
        )

    call_kwargs = mock_client.chat.completions.create.call_args.kwargs
    assert "reasoning_effort" not in call_kwargs


@pytest.mark.asyncio
async def test_extract_structured_vision_empty_images_raises():
    """Verify empty images list raises ValueError."""
    from srx_lib_llm.vision import extract_structured_vision

    with pytest.raises(ValueError, match="At least one image"):
        await extract_structured_vision(
            images=[],
            instruction="Extract.",
            json_schema=SAMPLE_SCHEMA,
        )


@pytest.mark.asyncio
async def test_extract_structured_vision_schema_preprocessing():
    """Verify additionalProperties:false and required enforcement."""
    from srx_lib_llm.vision import extract_structured_vision

    mock_client = MagicMock()
    mock_client.chat.completions.create.return_value = _mock_completion(SAMPLE_RESPONSE)

    with (
        patch("srx_lib_llm.vision.is_azure_openai", return_value=True),
        patch("srx_lib_llm.vision.get_effective_model", return_value="gpt-5-mini"),
        patch("srx_lib_llm.vision.create_openai_client", return_value=mock_client),
    ):
        await extract_structured_vision(
            images=[SAMPLE_IMAGE_URL],
            instruction="Extract.",
            json_schema=SAMPLE_SCHEMA,
        )

    call_kwargs = mock_client.chat.completions.create.call_args.kwargs
    schema = call_kwargs["response_format"]["json_schema"]["schema"]

    # All properties should be required (strict mode)
    assert set(schema["required"]) == {"item_name", "amount", "currency"}
    # additionalProperties should be false
    assert schema.get("additionalProperties") is False


@pytest.mark.asyncio
async def test_extract_structured_vision_custom_system():
    """Verify custom system message overrides default."""
    from srx_lib_llm.vision import extract_structured_vision

    mock_client = MagicMock()
    mock_client.chat.completions.create.return_value = _mock_completion(SAMPLE_RESPONSE)
    custom_sys = "You are a receipt OCR expert."

    with (
        patch("srx_lib_llm.vision.is_azure_openai", return_value=False),
        patch("srx_lib_llm.vision.get_effective_model", return_value="gpt-4.1-mini"),
        patch("srx_lib_llm.vision.get_openai_client_kwargs", return_value={}),
        patch("openai.OpenAI", return_value=mock_client),
    ):
        await extract_structured_vision(
            images=[SAMPLE_IMAGE_URL],
            instruction="Extract.",
            json_schema=SAMPLE_SCHEMA,
            system=custom_sys,
        )

    call_kwargs = mock_client.chat.completions.create.call_args.kwargs
    assert call_kwargs["messages"][0]["content"] == custom_sys


# ---------------------------------------------------------------------------
# Integration tests (requires Azure OpenAI credentials + test images)
# ---------------------------------------------------------------------------

@skip_no_azure
@pytest.mark.integration
@pytest.mark.asyncio
async def test_extract_structured_vision_real_azure():
    """Integration test: real Azure OpenAI call with a receipt image."""
    from srx_lib_llm.vision import extract_structured_vision

    # Use a test image from grab-receipt-manager
    img_path = _TEST_IMAGE_DIR / "type-b" / "trip_detail_cash.jpeg"
    if not img_path.exists():
        pytest.skip(f"Test image not found: {img_path}")

    img_bytes = img_path.read_bytes()
    data_uri = "data:image/jpeg;base64," + base64.b64encode(img_bytes).decode()

    schema = {
        "type": "object",
        "properties": {
            "trip_amount": {"type": "number", "description": "Total trip amount"},
            "payment_method": {
                "type": "string",
                "description": "Payment method: cash or cashless",
            },
            "confidence": {
                "type": "number",
                "description": "Extraction confidence 0-1",
            },
        },
        "required": ["trip_amount", "payment_method", "confidence"],
    }

    result = await extract_structured_vision(
        images=[data_uri],
        instruction="Extract the trip amount and payment method from this Grab receipt.",
        json_schema=schema,
        schema_name="TripReceipt",
        detail="high",
    )

    assert hasattr(result, "trip_amount")
    assert hasattr(result, "payment_method")
    assert result.trip_amount > 0
    assert result.payment_method in ("cash", "cashless")
