# Copyright (C) 2025 Embedl AB

"""Benchmarking of models using the Embedl device cloud."""

import tempfile
from pathlib import Path

from embedl_hub.core.benchmark.abc import Benchmarker, BenchmarkResult
from embedl_hub.core.utils.tracking_utils import (
    log_execution_detail,
    log_model_summary,
)
from embedl_hub.tracking import (
    log_param,
    submit_benchmark_job,
    validate_device,
)

RUNTIME = "TensorFlow Lite"


class EmbedlBenchmarker(Benchmarker):
    """Benchmarker that uses the Embedl device cloud."""

    supported_input_model_formats = {".tflite"}

    def __init__(self, device: str, artifacts_dir: Path | None = None):
        self._device = device
        self._artifacts_dir = artifacts_dir

        self._validate_device()

    def _validate_device(self) -> None:
        """Validate that the specified device is supported."""
        return validate_device(self._device)

    def _benchmark(self, model_path: Path) -> BenchmarkResult:
        """Benchmark a model on the Embedl device cloud and track results.

        Return a tuple of (summary_dict, execution_detail).
        """
        summary, _ = benchmark_model(
            model_path=model_path,
            device=self._device,
            artifacts_dir=self._artifacts_dir,
        )
        return BenchmarkResult(device=self._device, summary=summary)


def benchmark_model(
    model_path: Path, device: str, artifacts_dir: Path | None = None
) -> tuple[dict, list[dict]]:
    """Benchmark a model on the Embedl device cloud and track results.

    Return a tuple of (summary_dict, execution_detail).
    """

    log_param("$device", device)
    log_param("$runtime", RUNTIME)

    job = submit_benchmark_job(model_path=model_path, device=device)

    with tempfile.TemporaryDirectory() as temp_dir:
        artifacts_dir = artifacts_dir or Path(temp_dir)

        result = job.download_results(artifacts_dir=artifacts_dir)

        summary = result.summary
        execution_detail = result.execution_detail

        log_model_summary(summary)
        log_execution_detail(execution_detail)

    return summary, execution_detail
