"""Integration tests for Azure OpenAI support.

Tests are skipped when Azure env vars are not set.
Load env from grab-receipt-manager/.env.local for credentials.
"""
from __future__ import annotations

import os
import pathlib

import pytest

# ---------------------------------------------------------------------------
# Load Azure env vars from .env.local if available
# ---------------------------------------------------------------------------
_PROJECT_DIR = pathlib.Path(__file__).resolve().parent.parent
_ENV_FILES = [
    pathlib.Path.home() / "Documents/srx/codes/HOME/grab-receipt-manager/.env.local",
    _PROJECT_DIR / ".env",
]

def _load_env():
    for env_file in _ENV_FILES:
        if not env_file.exists():
            continue
        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"))
OPENAI_AVAILABLE = bool(os.getenv("OPENAI_API_KEY"))
BOTH_AVAILABLE = AZURE_AVAILABLE and OPENAI_AVAILABLE
skip_no_azure = pytest.mark.skipif(not AZURE_AVAILABLE, reason="Azure OpenAI env vars not set")
skip_no_both = pytest.mark.skipif(not BOTH_AVAILABLE, reason="Need both Azure and OpenAI credentials")


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------

@skip_no_azure
@pytest.mark.asyncio
async def test_responses_chat():
    from srx_lib_llm import responses_chat
    result = await responses_chat("Say hello in one word.")
    assert isinstance(result, str)
    assert len(result) > 0


@skip_no_azure
@pytest.mark.asyncio
async def test_responses_chat_with_cache():
    from srx_lib_llm import responses_chat
    result = await responses_chat("Say hello in one word.", cache=True)
    assert isinstance(result, str)
    assert len(result) > 0


@skip_no_azure
@pytest.mark.asyncio
async def test_extract_structured():
    from srx_lib_llm import extract_structured

    schema = {
        "type": "object",
        "properties": {
            "name": {"type": "string", "description": "Person name"},
            "age": {"type": "integer", "description": "Person age"},
        },
        "required": ["name", "age"],
    }
    result = await extract_structured(
        text="John is 25 years old.",
        json_schema=schema,
        schema_name="Person",
    )
    assert hasattr(result, "name")
    assert hasattr(result, "age")


@skip_no_azure
@pytest.mark.asyncio
async def test_extract_structured_with_model_override():
    """model param should be ignored on Azure; deployment is used instead."""
    from srx_lib_llm import extract_structured

    schema = {
        "type": "object",
        "properties": {
            "color": {"type": "string", "description": "Color name"},
        },
        "required": ["color"],
    }
    result = await extract_structured(
        text="The sky is blue.",
        json_schema=schema,
        schema_name="Color",
        model="gpt-4.1-mini",  # should be overridden by Azure deployment
    )
    assert hasattr(result, "color")


@skip_no_azure
@pytest.mark.asyncio
async def test_extract_structured_gpt51():
    """GPT-5 path via Azure chat completions."""
    from srx_lib_llm import extract_structured

    schema = {
        "type": "object",
        "properties": {
            "city": {"type": "string", "description": "City name"},
        },
        "required": ["city"],
    }
    # On Azure the deployment name starts with gpt-5, so extract_structured
    # will route to extract_structured_gpt51 automatically
    result = await extract_structured(
        text="I live in Jakarta.",
        json_schema=schema,
        schema_name="Location",
    )
    assert hasattr(result, "city")


@skip_no_azure
@pytest.mark.asyncio
async def test_extract_structured_gpt51_reasoning_none():
    from srx_lib_llm import extract_structured

    schema = {
        "type": "object",
        "properties": {
            "animal": {"type": "string", "description": "Animal name"},
        },
        "required": ["animal"],
    }
    result = await extract_structured(
        text="The cat sat on the mat.",
        json_schema=schema,
        schema_name="Animal",
        reasoning_effort="none",
    )
    assert hasattr(result, "animal")


def test_create_dynamic_schema():
    """Schema utilities work without API call."""
    from srx_lib_llm import create_dynamic_schema

    schema = {
        "type": "object",
        "properties": {
            "x": {"type": "string", "description": "X"},
        },
        "required": ["x"],
    }
    model_cls = create_dynamic_schema("TestSchema", schema)
    assert model_cls is not None
    assert "x" in model_cls.model_fields


def test_preprocess_json_schema():
    from srx_lib_llm import preprocess_json_schema

    schema = {
        "type": "object",
        "properties": {
            "a": {"type": "string"},
        },
    }
    result = preprocess_json_schema(schema)
    assert result["additionalProperties"] is False


@skip_no_azure
def test_batch_service_raises_on_azure():
    from srx_lib_llm import OpenAIBatchService

    with pytest.raises(NotImplementedError, match="Azure OpenAI does not support the Batch API"):
        OpenAIBatchService()


def test_is_azure_detection():
    from srx_lib_llm import is_azure_openai

    original = os.environ.get("AZURE_OPENAI_ENDPOINT")
    try:
        os.environ["AZURE_OPENAI_ENDPOINT"] = "https://example.openai.azure.com/"
        assert is_azure_openai() is True

        del os.environ["AZURE_OPENAI_ENDPOINT"]
        assert is_azure_openai() is False
    finally:
        if original:
            os.environ["AZURE_OPENAI_ENDPOINT"] = original
        elif "AZURE_OPENAI_ENDPOINT" in os.environ:
            del os.environ["AZURE_OPENAI_ENDPOINT"]


# ---------------------------------------------------------------------------
# Shape parity tests: Azure vs OpenAI must return identical structure
# ---------------------------------------------------------------------------

_PERSON_SCHEMA = {
    "type": "object",
    "properties": {
        "name": {"type": "string", "description": "Person name"},
        "age": {"type": "integer", "description": "Person age"},
    },
    "required": ["name", "age"],
}

_AZURE_ENV_KEYS = [
    "AZURE_OPENAI_ENDPOINT",
    "AZURE_OPENAI_API_KEY",
    "AZURE_OPENAI_DEPLOYMENT_NAME",
    "AZURE_OPENAI_API_VERSION",
]


def _save_azure_env():
    return {k: os.environ.get(k) for k in _AZURE_ENV_KEYS}


def _clear_azure_env():
    for k in _AZURE_ENV_KEYS:
        os.environ.pop(k, None)


def _restore_azure_env(saved):
    for k, v in saved.items():
        if v is not None:
            os.environ[k] = v
        else:
            os.environ.pop(k, None)


@skip_no_both
@pytest.mark.asyncio
async def test_shape_parity_responses_chat():
    """responses_chat must return the same type (str) from both providers."""
    from srx_lib_llm import responses_chat

    # Azure call
    azure_result = await responses_chat("Say 'parity test' and nothing else.")
    assert isinstance(azure_result, str) and len(azure_result) > 0

    # OpenAI call — temporarily clear Azure env
    saved = _save_azure_env()
    _clear_azure_env()
    try:
        openai_result = await responses_chat("Say 'parity test' and nothing else.")
    finally:
        _restore_azure_env(saved)

    assert isinstance(openai_result, str) and len(openai_result) > 0
    assert type(azure_result) is type(openai_result)


@skip_no_both
@pytest.mark.asyncio
async def test_shape_parity_extract_structured():
    """extract_structured must return identical field names from both providers."""
    from srx_lib_llm import extract_structured

    text = "Alice is 30 years old."

    # Azure
    azure_result = await extract_structured(
        text=text, json_schema=_PERSON_SCHEMA, schema_name="Person",
    )

    # OpenAI
    saved = _save_azure_env()
    _clear_azure_env()
    try:
        openai_result = await extract_structured(
            text=text, json_schema=_PERSON_SCHEMA, schema_name="Person",
        )
    finally:
        _restore_azure_env(saved)

    # Both must have same fields
    azure_fields = set(azure_result.model_dump().keys())
    openai_fields = set(openai_result.model_dump().keys())
    assert azure_fields == openai_fields, (
        f"Field mismatch: Azure={azure_fields}, OpenAI={openai_fields}"
    )

    # Both must have correct types
    azure_d = azure_result.model_dump()
    openai_d = openai_result.model_dump()
    for key in azure_fields:
        assert type(azure_d[key]) is type(openai_d[key]), (
            f"Type mismatch for '{key}': Azure={type(azure_d[key])}, OpenAI={type(openai_d[key])}"
        )


@skip_no_both
@pytest.mark.asyncio
async def test_shape_parity_extract_structured_nested():
    """Nested schema must produce identical structure from both providers."""
    from srx_lib_llm import extract_structured

    nested_schema = {
        "type": "object",
        "properties": {
            "company": {"type": "string", "description": "Company name"},
            "employees": {
                "type": "array",
                "description": "Employee list",
                "items": {
                    "type": "object",
                    "properties": {
                        "name": {"type": "string", "description": "Name"},
                        "role": {"type": "string", "description": "Role"},
                    },
                    "required": ["name", "role"],
                },
            },
        },
        "required": ["company", "employees"],
    }
    text = "Acme Corp has two employees: Bob the engineer and Carol the designer."

    # Azure
    azure_result = await extract_structured(
        text=text, json_schema=nested_schema, schema_name="Company",
    )

    # OpenAI
    saved = _save_azure_env()
    _clear_azure_env()
    try:
        openai_result = await extract_structured(
            text=text, json_schema=nested_schema, schema_name="Company",
        )
    finally:
        _restore_azure_env(saved)

    azure_d = azure_result.model_dump()
    openai_d = openai_result.model_dump()

    # Same top-level keys
    assert set(azure_d.keys()) == set(openai_d.keys())

    # employees is a list in both
    assert isinstance(azure_d["employees"], list)
    assert isinstance(openai_d["employees"], list)

    # Each employee item has same keys
    if azure_d["employees"] and openai_d["employees"]:
        assert set(azure_d["employees"][0].keys()) == set(openai_d["employees"][0].keys())
