# SPDX-License-Identifier: Apache-2.0
"""Tests for PaddleOCRv5 ONNX models.

This test downloads PaddleOCRv5 ONNX models from GitHub
and verifies that onnx2fx can convert and run them correctly.

Models from: https://github.com/Kazuhito00/PaddleOCRv5-ONNX-Sample
"""

import urllib.request
from pathlib import Path

import numpy as np
import onnx
import onnxruntime as ort
import pytest
import torch

from conftest import run_onnx_test

# Base URL for model downloads
MODEL_BASE_URL = "https://github.com/Kazuhito00/PaddleOCRv5-ONNX-Sample/raw/main"

# Model paths (relative to repository root)
MODEL_FILES = {
    "mobile_det": "ppocr_onnx/model/det_model/PP-OCRv5_mobile_det_infer.onnx",
    "mobile_rec": "ppocr_onnx/model/rec_model/PP-OCRv5_mobile_rec_infer.onnx",
    "server_det": "ppocr_onnx/model/det_model/PP-OCRv5_server_det_infer.onnx",
    "server_rec": "ppocr_onnx/model/rec_model/PP-OCRv5_server_rec_infer.onnx",
}

# Local cache filenames
MODEL_CACHE_NAMES = {
    "mobile_det": "PP-OCRv5_mobile_det_infer.onnx",
    "mobile_rec": "PP-OCRv5_mobile_rec_infer.onnx",
    "server_det": "PP-OCRv5_server_det_infer.onnx",
    "server_rec": "PP-OCRv5_server_rec_infer.onnx",
}

# Sample image URL
SAMPLE_IMAGE_URL = f"{MODEL_BASE_URL}/sample.jpg"
SAMPLE_IMAGE_NAME = "sample.jpg"

# Cache directory for downloaded models
CACHE_DIR = Path(__file__).parent / ".model_cache"


def download_model(model_key: str) -> Path:
    """Download PaddleOCRv5 model from GitHub.

    Args:
        model_key: Key for the model (mobile_det, mobile_rec, server_det, server_rec)

    Returns:
        Path to downloaded ONNX model file.
    """
    model_path_in_repo = MODEL_FILES[model_key]
    cache_filename = MODEL_CACHE_NAMES[model_key]
    url = f"{MODEL_BASE_URL}/{model_path_in_repo}"

    # Create cache directory if it doesn't exist
    CACHE_DIR.mkdir(parents=True, exist_ok=True)

    model_path = CACHE_DIR / cache_filename

    # Download if not already cached
    if not model_path.exists():
        print(f"Downloading {cache_filename} from {url}...")
        urllib.request.urlretrieve(url, model_path)
        print(f"Downloaded to {model_path}")

    return model_path


def download_sample_image() -> Path:
    """Download sample.jpg from GitHub.

    Returns:
        Path to downloaded image file.
    """
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    image_path = CACHE_DIR / SAMPLE_IMAGE_NAME

    if not image_path.exists():
        print(f"Downloading {SAMPLE_IMAGE_NAME} from {SAMPLE_IMAGE_URL}...")
        urllib.request.urlretrieve(SAMPLE_IMAGE_URL, image_path)
        print(f"Downloaded to {image_path}")

    return image_path


def load_and_preprocess_image(
    image_path: Path, target_height: int, target_width: int
) -> torch.Tensor:
    """Load and preprocess image for OCR model input.

    Args:
        image_path: Path to the image file.
        target_height: Target height for resizing.
        target_width: Target width for resizing.

    Returns:
        Preprocessed image tensor of shape (1, 3, H, W).
    """
    from PIL import Image

    # Load image
    img = Image.open(image_path).convert("RGB")

    # Resize to target size
    img = img.resize((target_width, target_height), Image.BILINEAR)

    # Convert to numpy array and normalize to [0, 1]
    img_array = np.array(img, dtype=np.float32) / 255.0

    # Normalize with mean and std (common OCR preprocessing)
    mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
    std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
    img_array = (img_array - mean) / std

    # HWC to CHW and add batch dimension
    img_array = img_array.transpose(2, 0, 1)
    img_tensor = torch.from_numpy(img_array).unsqueeze(0)

    return img_tensor


# --- Fixtures ---


@pytest.fixture(scope="module")
def mobile_det_model_path():
    """Download and cache the mobile detection model."""
    return download_model("mobile_det")


@pytest.fixture(scope="module")
def mobile_rec_model_path():
    """Download and cache the mobile recognition model."""
    return download_model("mobile_rec")


@pytest.fixture(scope="module")
def server_det_model_path():
    """Download and cache the server detection model."""
    return download_model("server_det")


@pytest.fixture(scope="module")
def server_rec_model_path():
    """Download and cache the server recognition model."""
    return download_model("server_rec")


@pytest.fixture(scope="module")
def mobile_det_onnx_model(mobile_det_model_path):
    """Load the mobile detection ONNX model."""
    return onnx.load(str(mobile_det_model_path))


@pytest.fixture(scope="module")
def mobile_rec_onnx_model(mobile_rec_model_path):
    """Load the mobile recognition ONNX model."""
    return onnx.load(str(mobile_rec_model_path))


@pytest.fixture(scope="module")
def server_det_onnx_model(server_det_model_path):
    """Load the server detection ONNX model."""
    return onnx.load(str(server_det_model_path))


@pytest.fixture(scope="module")
def server_rec_onnx_model(server_rec_model_path):
    """Load the server recognition ONNX model."""
    return onnx.load(str(server_rec_model_path))


@pytest.fixture(scope="module")
def sample_image_path():
    """Download and cache the sample image."""
    return download_sample_image()


class TestPaddleOCRv5Inference:
    """Test PaddleOCRv5 model inference comparing FX with ONNX Runtime."""

    @pytest.mark.slow
    def test_mobile_det_inference(
        self, mobile_det_model_path, mobile_det_onnx_model, sample_image_path
    ):
        """Test mobile detection model inference matches ONNX Runtime."""
        # Load and preprocess sample image
        test_input = load_and_preprocess_image(sample_image_path, 640, 640)
        np_input = test_input.numpy()

        # Run ONNX Runtime
        ort_session = ort.InferenceSession(
            str(mobile_det_model_path), providers=["CPUExecutionProvider"]
        )
        input_name = ort_session.get_inputs()[0].name
        ort_outputs = ort_session.run(None, {input_name: np_input})

        expected = torch.from_numpy(ort_outputs[0].copy())
        run_onnx_test(
            mobile_det_onnx_model,
            test_input,
            expected,
            rtol=0.1,
            atol=1.0,
            output_transform=lambda out: out[0]
            if isinstance(out, (tuple, list))
            else out,
        )

    @pytest.mark.slow
    def test_mobile_rec_inference(
        self, mobile_rec_model_path, mobile_rec_onnx_model, sample_image_path
    ):
        """Test mobile recognition model inference matches ONNX Runtime."""
        # Load and preprocess sample image (recognition uses different size)
        test_input = load_and_preprocess_image(sample_image_path, 48, 320)
        np_input = test_input.numpy()

        # Run ONNX Runtime
        ort_session = ort.InferenceSession(
            str(mobile_rec_model_path), providers=["CPUExecutionProvider"]
        )
        input_name = ort_session.get_inputs()[0].name
        ort_outputs = ort_session.run(None, {input_name: np_input})

        expected = torch.from_numpy(ort_outputs[0].copy())
        run_onnx_test(
            mobile_rec_onnx_model,
            test_input,
            expected,
            rtol=0.05,
            atol=0.1,
            output_transform=lambda out: out[0]
            if isinstance(out, (tuple, list))
            else out,
        )

    @pytest.mark.slow
    def test_server_det_inference(
        self, server_det_model_path, server_det_onnx_model, sample_image_path
    ):
        """Test server detection model inference matches ONNX Runtime."""
        # Load and preprocess sample image
        test_input = load_and_preprocess_image(sample_image_path, 640, 640)
        np_input = test_input.numpy()

        # Run ONNX Runtime
        ort_session = ort.InferenceSession(
            str(server_det_model_path), providers=["CPUExecutionProvider"]
        )
        input_name = ort_session.get_inputs()[0].name
        ort_outputs = ort_session.run(None, {input_name: np_input})

        expected = torch.from_numpy(ort_outputs[0].copy())
        run_onnx_test(
            server_det_onnx_model,
            test_input,
            expected,
            rtol=0.05,
            atol=0.1,
            output_transform=lambda out: out[0]
            if isinstance(out, (tuple, list))
            else out,
        )

    @pytest.mark.slow
    def test_server_rec_inference(
        self, server_rec_model_path, server_rec_onnx_model, sample_image_path
    ):
        """Test server recognition model inference matches ONNX Runtime."""
        # Load and preprocess sample image (recognition uses different size)
        test_input = load_and_preprocess_image(sample_image_path, 48, 320)
        np_input = test_input.numpy()

        # Run ONNX Runtime
        ort_session = ort.InferenceSession(
            str(server_rec_model_path), providers=["CPUExecutionProvider"]
        )
        input_name = ort_session.get_inputs()[0].name
        ort_outputs = ort_session.run(None, {input_name: np_input})

        expected = torch.from_numpy(ort_outputs[0].copy())
        run_onnx_test(
            server_rec_onnx_model,
            test_input,
            expected,
            rtol=0.05,
            atol=0.1,
            output_transform=lambda out: out[0]
            if isinstance(out, (tuple, list))
            else out,
        )
