"""OpenVINO model exporter."""

from __future__ import annotations

import logging
from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import torch

from matrice_export.formats.base import BaseExporter

logger = logging.getLogger(__name__)


class OpenVinoExporter(BaseExporter):
    """Export a PyTorch model to OpenVINO Intermediate Representation (IR).

    The pipeline works in two stages:

    1. An ONNX file must already exist (or ``onnx_path`` is supplied).
       If neither is present the exporter will call
       :class:`~matrice_export.formats.onnx_export.OnnxExporter` internally
       to produce the prerequisite ONNX model.
    2. The ONNX graph is converted to OpenVINO IR via
       ``openvino.tools.mo.convert_model`` (FP32 / FP16) or quantised to
       INT8 with ``nncf.quantize``.
    """

    @property
    def format_name(self) -> str:
        return "openvino"

    @property
    def suffix(self) -> str:
        return "_openvino_model"

    def export(
        self,
        model: torch.nn.Module,
        sample_input: torch.Tensor,
        output_dir: str | Path,
        file_stem: str = "model",
        **kwargs,
    ) -> str:
        """Export a PyTorch model to OpenVINO IR.

        Args:
            model: PyTorch model in eval mode (used only when an ONNX
                intermediate must be generated).
            sample_input: Sample input tensor (BCHW).
            output_dir: Directory to write the exported directory into.
            file_stem: Base filename without extension.
            **kwargs:
                half (bool): Compress weights to FP16.  Defaults to ``False``.
                int8 (bool): Apply NNCF post-training INT8 quantisation.
                    Defaults to ``False``.
                data_dir (str | Path | None): Path to a calibration image
                    directory (``ImageFolder`` layout) required when
                    ``int8=True``.
                onnx_path (str | Path | None): Path to an existing ONNX model.
                    When ``None`` the exporter will create one in *output_dir*.
                metadata (dict | None): Extra metadata to persist as a YAML
                    sidecar file.
                stride (int | None): Model stride value for metadata.
                names (dict | list | None): Class names for metadata.

        Returns:
            Absolute path to the exported OpenVINO model directory.
        """
        import torch

        try:
            import openvino.runtime as ov  # noqa: F811
            from openvino.tools import mo
        except ImportError as exc:
            raise ImportError(
                "The 'openvino-dev' package (>= 2023.0) is required for "
                "OpenVINO export.  Install it with:  pip install openvino-dev>=2023.0"
            ) from exc

        half: bool = kwargs.get("half", False)
        int8: bool = kwargs.get("int8", False)
        data_dir: str | Path | None = kwargs.get("data_dir", None)
        onnx_path: str | Path | None = kwargs.get("onnx_path", None)
        metadata: dict | None = kwargs.get("metadata", None)

        output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)

        # Target directory for the OpenVINO model
        ov_dir = output_dir / f"{file_stem}{self.suffix}"
        ov_dir.mkdir(parents=True, exist_ok=True)
        f_ov = str(ov_dir / f"{file_stem}.xml")

        # ------------------------------------------------------------------
        # Ensure we have an ONNX file to convert
        # ------------------------------------------------------------------
        if onnx_path is None:
            onnx_path = output_dir / f"{file_stem}.onnx"
            if not Path(onnx_path).exists():
                logger.info(
                    "OpenVINO export: ONNX prerequisite not found, "
                    "running OnnxExporter first ..."
                )
                from matrice_export.formats.onnx_export import OnnxExporter

                onnx_exporter = OnnxExporter()
                onnx_path = onnx_exporter.export(
                    model,
                    sample_input,
                    output_dir,
                    file_stem=file_stem,
                    **{k: v for k, v in kwargs.items() if k in ("opset", "dynamic", "simplify", "stride", "names")},
                )
        onnx_path = str(onnx_path)

        logger.info("OpenVINO export: converting ONNX -> IR ...")

        # ------------------------------------------------------------------
        # INT8 quantisation path
        # ------------------------------------------------------------------
        if int8:
            try:
                import nncf
            except ImportError as exc:
                raise ImportError(
                    "The 'nncf' package (>= 2.4.0) is required for INT8 "
                    "quantisation.  Install it with:  pip install nncf>=2.4.0"
                ) from exc

            try:
                import numpy as np
                from openvino.runtime import Core
                from torchvision import datasets, transforms
            except ImportError as exc:
                raise ImportError(
                    "torchvision, numpy, and openvino.runtime are required "
                    "for INT8 calibration."
                ) from exc

            if data_dir is None:
                raise ValueError(
                    "data_dir must be provided when int8=True so that "
                    "calibration data can be loaded."
                )

            transform = transforms.Compose([
                transforms.Resize(256),
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                transforms.Normalize(
                    mean=[0.485, 0.456, 0.406],
                    std=[0.229, 0.224, 0.225],
                ),
            ])

            dataset = datasets.ImageFolder(str(data_dir), transform=transform)
            dataloader = torch.utils.data.DataLoader(
                dataset, batch_size=1, shuffle=False,
            )

            core = Core()
            onnx_model = core.read_model(onnx_path)

            def _prepare_input(image: np.ndarray) -> np.ndarray:
                tensor = image.astype(np.float32)
                tensor /= 255.0
                if tensor.ndim == 3:
                    tensor = np.expand_dims(tensor, 0)
                return tensor

            def _transform_fn(data_item):
                img = data_item[0].numpy()
                return _prepare_input(img)

            calibration_dataset = nncf.Dataset(dataloader, _transform_fn)
            ov_model = nncf.quantize(
                onnx_model,
                calibration_dataset,
                preset=nncf.QuantizationPreset.MIXED,
            )
            logger.info("OpenVINO export: INT8 quantisation complete.")
        else:
            # ------------------------------------------------------------------
            # FP32 / FP16 conversion
            # ------------------------------------------------------------------
            ov_model = mo.convert_model(
                onnx_path,
                model_name=file_stem,
                framework="onnx",
                compress_to_fp16=half,
            )

        # ------------------------------------------------------------------
        # Serialise the IR model
        # ------------------------------------------------------------------
        ov.serialize(ov_model, f_ov)
        logger.info("OpenVINO export: serialised to %s", f_ov)

        # ------------------------------------------------------------------
        # Optional metadata sidecar (YAML)
        # ------------------------------------------------------------------
        if metadata is not None:
            try:
                import yaml
            except ImportError:
                logger.warning(
                    "OpenVINO export: PyYAML not installed -- skipping "
                    "metadata YAML sidecar."
                )
            else:
                yaml_path = ov_dir / f"{file_stem}.yaml"
                with open(yaml_path, "w", encoding="utf-8") as fh:
                    yaml.safe_dump(metadata, fh, default_flow_style=False)
                logger.info("OpenVINO export: metadata saved to %s", yaml_path)

        logger.info("OpenVINO export: saved to %s", ov_dir)
        return str(ov_dir)
