from pathlib import Path
from typing import Any, TypedDict

import qai_hub.hub as hub
from qai_hub.client import CompileJob, Device, InferenceJob, Model, ProfileJob, QuantizeDtype, QuantizeJob


class ModelJobResult(TypedDict):
    job: CompileJob | QuantizeJob
    job_id: str
    model: Model
    model_id: str


class InferenceJobResult(TypedDict):
    job: InferenceJob
    job_id: str
    outputs: Any


class ProfileJobResult(TypedDict):
    job: ProfileJob
    job_id: str


def _dataset_entries(inputs: dict[str, Any]) -> dict[str, list[Any]]:
    return {name: value if isinstance(value, list) else [value] for name, value in inputs.items()}


def submit_compile_job(
    model: Model,
    device: Device,
    input_specs: dict[str, tuple[tuple[int, ...], str]],
    target_runtime: str,
    options: str | None = None,
    job_name: str | None = None,
) -> ModelJobResult:
    compile_options = f"--target_runtime {target_runtime}"
    if options:
        compile_options = f"{compile_options} {options}"

    job = hub.submit_compile_job(
        model=model,
        device=device,
        name=job_name,
        input_specs=input_specs,
        options=compile_options,
    )
    target_model = job.get_target_model()
    if target_model is None:
        raise RuntimeError(f"Compile job {job.job_id} did not produce a target model.")
    return {"job": job, "job_id": str(job.job_id), "model": target_model, "model_id": str(target_model.model_id)}


def submit_inference_job(
    model: Model,
    device: Device,
    inputs: dict[str, Any],
    output_dir: str | Path,
    job_name: str | None = None,
) -> InferenceJobResult:
    job = hub.submit_inference_job(
        model=model,
        device=device,
        inputs=_dataset_entries(inputs),
        name=job_name,
    )
    out = Path(output_dir)
    out.mkdir(parents=True, exist_ok=True)
    data = job.download_output_data(str(out))
    return {"job": job, "job_id": str(job.job_id), "outputs": data}


def submit_profile_job(
    model: Model,
    device: Device,
    options: str | None = None,
    job_name: str | None = None,
) -> ProfileJobResult:
    job = hub.submit_profile_job(
        model=model,
        device=device,
        name=job_name,
        options=options or "",
    )
    return {"job": job, "job_id": str(job.job_id)}


def submit_quantize_job(
    model: Model,
    calibration_data: dict[str, Any],
    options: str | None = None,
    job_name: str | None = None,
) -> ModelJobResult:
    job = hub.submit_quantize_job(
        model=model,
        calibration_data=_dataset_entries(calibration_data),
        weights_dtype=QuantizeDtype.INT8,
        activations_dtype=QuantizeDtype.INT8,
        name=job_name,
        options=options or "",
    )
    target_model = job.get_target_model()
    if target_model is None:
        raise RuntimeError(f"Quantize job {job.job_id} did not produce a target model.")
    return {"job": job, "job_id": str(job.job_id), "model": target_model, "model_id": str(target_model.model_id)}


def download_model(model_id: str, output_path: str | Path) -> str:
    dest = Path(output_path)
    dest.parent.mkdir(parents=True, exist_ok=True)
    model = hub.get_model(model_id)
    result = model.download(str(dest))
    return str(result or dest)