"""TensorFlow Lite (``.tflite``) exporter.

Converts a TF SavedModel to TFLite format with FP16 quantisation by default
and optional INT8 quantisation using a representative dataset.
"""

from __future__ import annotations

import contextlib
import logging
from pathlib import Path
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    import torch

from matrice_export.formats.base import BaseExporter

logger = logging.getLogger(__name__)


class TFLiteExporter(BaseExporter):
    """Export a model to TensorFlow Lite (``.tflite``)."""

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

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

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------
    def export(
        self,
        model: torch.nn.Module,
        sample_input: torch.Tensor,
        output_dir: str | Path,
        file_stem: str = "model",
        **kwargs: Any,
    ) -> str:
        """Export to TFLite.

        Keyword Args:
            saved_model_path (str | Path | None): Path to an existing TF
                SavedModel directory.  When *None*, one is produced via
                :class:`TFSavedModelExporter`.
            int8 (bool): If ``True``, apply full INT8 quantisation with a
                representative dataset (requires ``representative_dataset``).
                Defaults to ``False`` (FP16 quantisation).
            representative_dataset (Iterator | None): A callable that yields
                representative input samples for INT8 calibration.
            nms (bool): Whether the model uses TF NMS ops (adds SELECT_TF_OPS).
            agnostic_nms (bool): Same as *nms* (adds SELECT_TF_OPS).
            metadata (dict | None): Optional metadata dict to embed in the
                TFLite file.  Requires ``tflite_support``.
            num_outputs (int): Number of model outputs (used when embedding
                metadata).

        Returns:
            Path to the exported ``.tflite`` file.
        """
        output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)

        int8: bool = kwargs.get("int8", False)
        nms: bool = kwargs.get("nms", False)
        agnostic_nms: bool = kwargs.get("agnostic_nms", False)
        representative_dataset = kwargs.get("representative_dataset")
        metadata: dict | None = kwargs.get("metadata")
        num_outputs: int = kwargs.get("num_outputs", 1)

        quantisation_tag = "int8" if int8 else "fp16"
        tflite_path = output_dir / f"{file_stem}-{quantisation_tag}{self.suffix}"

        saved_model_path: str | Path | None = kwargs.get("saved_model_path")

        # Ensure a SavedModel is available.
        if saved_model_path is None:
            from matrice_export.formats.tf_saved import TFSavedModelExporter

            logger.info("No saved_model_path supplied — creating TF SavedModel first ...")
            saved_model_path = TFSavedModelExporter().export(
                model, sample_input, output_dir, file_stem, **kwargs
            )

        saved_model_path = Path(saved_model_path)
        if not saved_model_path.exists():
            raise FileNotFoundError(
                f"TF SavedModel directory not found: {saved_model_path}"
            )

        # --- Convert to TFLite ----------------------------------------
        try:
            import tensorflow as tf
        except ImportError as exc:
            raise ImportError(
                "TFLite export requires TensorFlow. "
                "Install it with:  pip install tensorflow"
            ) from exc

        logger.info(
            "Converting TF SavedModel (%s) to TFLite (%s) [%s] ...",
            saved_model_path,
            tflite_path,
            quantisation_tag.upper(),
        )

        converter = tf.lite.TFLiteConverter.from_saved_model(str(saved_model_path))
        converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
        converter.optimizations = [tf.lite.Optimize.DEFAULT]

        if int8:
            if representative_dataset is None:
                raise ValueError(
                    "INT8 quantisation requires a representative dataset. "
                    "Pass it via representative_dataset=<callable>."
                )
            converter.representative_dataset = representative_dataset
            converter.target_spec.supported_ops = [
                tf.lite.OpsSet.TFLITE_BUILTINS_INT8,
            ]
            converter.target_spec.supported_types = []
            converter.inference_input_type = tf.uint8
            converter.inference_output_type = tf.uint8
            converter.experimental_new_quantizer = True
        else:
            # Default: FP16 quantisation
            converter.target_spec.supported_types = [tf.float16]

        if nms or agnostic_nms:
            converter.target_spec.supported_ops.append(
                tf.lite.OpsSet.SELECT_TF_OPS
            )

        tflite_model = converter.convert()
        tflite_path.write_bytes(tflite_model)

        # Optionally embed metadata
        if metadata is not None:
            self.add_tflite_metadata(str(tflite_path), metadata, num_outputs)

        logger.info("TFLite export complete: %s", tflite_path)
        return str(tflite_path)

    # ------------------------------------------------------------------
    # Metadata helper
    # ------------------------------------------------------------------
    @staticmethod
    def add_tflite_metadata(
        file: str,
        metadata: dict,
        num_outputs: int,
    ) -> None:
        """Embed metadata into an existing ``.tflite`` file.

        Uses the ``tflite_support`` package.  If the package is not installed
        the call is silently skipped.

        Args:
            file: Path to the ``.tflite`` file.
            metadata: Arbitrary metadata dictionary (serialised as text).
            num_outputs: Number of model output tensors.
        """
        with contextlib.suppress(ImportError):
            from tflite_support import flatbuffers  # type: ignore[import-untyped]
            from tflite_support import metadata as _metadata  # type: ignore[import-untyped]
            from tflite_support import metadata_schema_py_generated as _metadata_fb  # type: ignore[import-untyped]

            tmp_file = Path("/tmp/meta.txt")
            tmp_file.write_text(str(metadata))

            model_meta = _metadata_fb.ModelMetadataT()
            label_file = _metadata_fb.AssociatedFileT()
            label_file.name = tmp_file.name
            model_meta.associatedFiles = [label_file]

            subgraph = _metadata_fb.SubGraphMetadataT()
            subgraph.inputTensorMetadata = [_metadata_fb.TensorMetadataT()]
            subgraph.outputTensorMetadata = [
                _metadata_fb.TensorMetadataT()
            ] * num_outputs
            model_meta.subgraphMetadata = [subgraph]

            b = flatbuffers.Builder(0)
            b.Finish(
                model_meta.Pack(b),
                _metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER,
            )
            metadata_buf = b.Output()

            populator = _metadata.MetadataPopulator.with_model_file(file)
            populator.load_metadata_buffer(metadata_buf)
            populator.load_associated_files([str(tmp_file)])
            populator.populate()
            tmp_file.unlink()

            logger.info("TFLite metadata embedded into %s", file)
