# Copyright (C) 2025 Embedl AB

"""
Utils for Qualcomm AI Hub integration.
"""

import shutil
import zipfile
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import Optional, Union

import onnx
import qai_hub as hub
import qai_hub.public_api_pb2 as hub_pb
import qai_hub.public_rest_api as hub_api
from qai_hub.client import Client, Model
from qai_hub.hub import _global_client

QAI_HUB_RUNTIME_NAMES = ["ONNX Runtime", "TensorFlow Lite", "QNN"]


def get_global_qai_hub_client() -> Client:
    """Get the global Qualcomm AI Hub API client."""
    return _global_client


def get_job_result(
    job_id: str, client_config: hub_api.ClientConfig
) -> hub_pb.JobResult:
    """Get the protobuf representation of a job result from Qualcomm AI Hub.

    Used to retrieve information about a job that is otherwise not exposed by
    the `qai_hub` package.
    """
    # pylint: disable-next=protected-access
    job_result: hub_pb.JobResult = hub.client._api_call(
        hub_api.get_job_results, client_config, job_id
    )
    return job_result


def parse_runtime_info(job_result: hub_pb.JobResult) -> str:
    """Extract the runtime name from job result protobuf.

    If no recognized runtime is found, an error will be raised.
    """

    runtime = None

    job_type = job_result.WhichOneof("result")

    job_runtime_names = []

    if job_type == "compile_job_result":
        job_runtime_names = [
            tool.name
            for tool in job_result.compile_job_result.compile_detail.tool_versions
        ]
    elif job_type == "profile_job_result":
        job_runtime_names = [
            runtime.name
            for runtime in job_result.profile_job_result.profile.runtime_config
        ]
    else:
        raise RuntimeError(f"Unrecognized job type: {job_type}.")

    for expected_runtime_name in QAI_HUB_RUNTIME_NAMES:
        if expected_runtime_name in job_runtime_names:
            runtime = expected_runtime_name
            break

    if runtime is None:
        raise RuntimeError(
            f"No recognized runtime in job result: {job_runtime_names}. "
            f"Expected one of: {QAI_HUB_RUNTIME_NAMES}."
        )

    return runtime


def unzip_if_zipped(model_path: Path, tmp_folder: Path) -> Path:
    """Unzip the model if it is a zip file and return the ONNX file path."""
    if model_path.suffix == ".zip":
        with zipfile.ZipFile(model_path, "r") as zip_ref:
            extract_path = tmp_folder / model_path.stem
            zip_ref.extractall(extract_path)
            subfolders = [f for f in extract_path.iterdir() if f.is_dir()]
            if len(subfolders) != 1:
                raise RuntimeError(
                    f"Expected exactly one folder in the zip, found: {len(subfolders)}"
                )
            only_folder = subfolders[0]
            onnx_files = list(only_folder.rglob("*.onnx"))
            if len(onnx_files) != 1:
                raise RuntimeError(
                    f"Expected exactly one .onnx file, found: {len(onnx_files)}"
                )
            return onnx_files[0]
    return model_path


def save_qai_hub_model(
    model: Model,
    output_file: Path | str | None = None,
) -> Path:
    """
    Save a Qualcomm AI Hub model to a local file.

    If the model is a zip archive, it will be extracted. If the model is small enough,
    it will be saved as a single ONNX file. If the model is too large, it will be saved
    in a folder with multiple files.

    If the model is not a zip archive, it will be saved directly.
    """

    if output_file is None:
        output_file = Path.cwd() / model.name
    else:
        output_file = Path(output_file)
    output_file.parent.mkdir(parents=True, exist_ok=True)

    with TemporaryDirectory() as tmpdir:
        tmp_file = Path(tmpdir) / "model"
        tmp_zip = model.download(str(tmp_file))
        model_file = unzip_if_zipped(Path(tmp_zip), Path(tmpdir))
        if model_file.suffix != ".onnx":
            shutil.copy2(model_file, output_file)
            return output_file
        onnx_model = onnx.load(model_file)
        try:
            onnx.save_model(onnx_model, output_file)
            return output_file
        except Exception:
            # Model too large to save as single file, save as folder
            output_folder = output_file.with_suffix("")
            output_folder.mkdir(parents=True, exist_ok=True)
            output_file = (output_folder / model.name).with_suffix(".onnx")
            data_file = output_file.with_suffix(".data")
            onnx.save_model(
                onnx_model,
                output_file,
                save_as_external_data=True,
                location=data_file.name,
            )
            return output_folder
