"""Edge TPU exporter (``_edgetpu.tflite``).

Compiles an INT8-quantised TFLite model for the Google Coral Edge TPU using
the ``edgetpu_compiler`` CLI.  Linux only.
"""

from __future__ import annotations

import logging
import platform
import subprocess
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__)

_HELP_URL = "https://coral.ai/docs/edgetpu/compiler/"


class EdgeTPUExporter(BaseExporter):
    """Compile an INT8 TFLite model for the Google Coral Edge TPU."""

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

    @property
    def suffix(self) -> str:
        return "_edgetpu.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:
        """Compile an INT8 TFLite model for the Edge TPU.

        This exporter is **Linux-only** and requires the
        ``edgetpu_compiler`` CLI.  If the compiler is not found it will
        attempt to install it automatically.

        Keyword Args:
            tflite_path (str | Path | None): Path to an existing INT8
                ``.tflite`` file.  When *None*, the model is first exported
                to INT8 TFLite via :class:`TFLiteExporter`.
            Any additional kwargs are forwarded to upstream exporters if
            the TFLite file needs to be created.

        Returns:
            Path to the compiled ``*_edgetpu.tflite`` file.
        """
        if platform.system() != "Linux":
            raise OSError(
                f"Edge TPU export is only supported on Linux. See {_HELP_URL}"
            )

        output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        edgetpu_path = output_dir / f"{file_stem}{self.suffix}"

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

        # Ensure an INT8 TFLite model is available.
        if tflite_path is None:
            from matrice_export.formats.tflite import TFLiteExporter

            logger.info(
                "No tflite_path supplied — creating INT8 TFLite model first ..."
            )
            tflite_kwargs = {**kwargs, "int8": True}
            tflite_path = TFLiteExporter().export(
                model, sample_input, output_dir, file_stem, **tflite_kwargs
            )

        tflite_path = Path(tflite_path)
        if not tflite_path.exists():
            raise FileNotFoundError(f"TFLite file not found: {tflite_path}")

        # Ensure the Edge TPU compiler is installed.
        self._ensure_compiler()

        # Get compiler version for logging.
        ver = (
            subprocess.run(
                ["edgetpu_compiler", "--version"],
                capture_output=True,
                check=True,
                timeout=30,
            )
            .stdout.decode()
            .strip()
            .split()[-1]
        )
        logger.info(
            "Compiling Edge TPU model with edgetpu_compiler %s ...", ver
        )

        # Compile.
        subprocess.run(
            [
                "edgetpu_compiler",
                "-s",
                "-d",
                "-k",
                "10",
                "--out_dir",
                str(output_dir),
                str(tflite_path),
            ],
            check=True,
            timeout=300,
        )

        # The compiler writes the output next to the source with an
        # ``_edgetpu`` infix.  Rename to our canonical location if needed.
        compiler_output = tflite_path.with_name(
            tflite_path.stem + "_edgetpu.tflite"
        )
        if compiler_output.exists() and compiler_output != edgetpu_path:
            compiler_output.rename(edgetpu_path)

        if not edgetpu_path.exists():
            raise RuntimeError(
                f"Edge TPU compilation did not produce {edgetpu_path}. "
                "Check edgetpu_compiler output for errors."
            )

        logger.info("Edge TPU export complete: %s", edgetpu_path)
        return str(edgetpu_path)

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------
    @staticmethod
    def _ensure_compiler() -> None:
        """Check for ``edgetpu_compiler`` and raise if missing."""
        try:
            subprocess.run(
                ["edgetpu_compiler", "--version"],
                capture_output=True,
                check=True,
                timeout=30,
            )
        except (FileNotFoundError, subprocess.CalledProcessError):
            raise RuntimeError(
                "edgetpu_compiler not found. Please install it manually:\n"
                f"  See {_HELP_URL}\n"
                "  On Debian/Ubuntu:\n"
                "    curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -\n"
                '    echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main"'
                " | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list\n"
                "    sudo apt-get update && sudo apt-get install edgetpu-compiler"
            )
