from __future__ import annotations

from dataclasses import dataclass

from models import HFModelSummary, ModelInspection


VLLM_LIKE_FRAMEWORKS = (
    "vllm",
    "sglang",
    "vllm-customized",
    "vllm-mlu",
    "vllm-016",
    "vllm_fix_tokenizer",
)


@dataclass(frozen=True)
class TaskSpec:
    task_type: str
    modality: str
    pipeline_tags: tuple[str, ...]
    priority: int


TASK_SPECS: tuple[TaskSpec, ...] = (
    TaskSpec("text-generation", "text", ("text-generation",), 10),
    TaskSpec(
        "visual-multi-modal",
        "multimodal",
        ("image-text-to-text", "visual-question-answering", "document-question-answering", "video-text-to-text"),
        20,
    ),
    TaskSpec("text-to-image-generation", "image", ("text-to-image", "image-to-image"), 30),
    TaskSpec("asr", "audio", ("automatic-speech-recognition",), 40),
    TaskSpec("question_answering", "text", ("question-answering",), 50),
    TaskSpec("feature_emb", "embedding", ("feature-extraction", "sentence-similarity"), 60),
    TaskSpec("vision_classification", "vision", ("image-classification", "zero-shot-image-classification"), 70),
    TaskSpec("text_classification", "text", ("text-classification", "zero-shot-classification"), 80),
    TaskSpec("reinforcement_learning", "text", ("reinforcement-learning",), 90),
)


TASK_SPEC_BY_TYPE = {task.task_type: task for task in TASK_SPECS}


def all_task_types() -> list[str]:
    return [task.task_type for task in TASK_SPECS]


def pipeline_tags_for_task_types(task_types: list[str]) -> list[str]:
    tags: list[str] = []
    for task_type in task_types:
        spec = TASK_SPEC_BY_TYPE[task_type]
        for tag in spec.pipeline_tags:
            if tag not in tags:
                tags.append(tag)
    return tags


def task_specs_for_model(model: HFModelSummary) -> list[TaskSpec]:
    pipeline_tag = (model.pipeline_tag or "").strip().lower()
    return [task for task in TASK_SPECS if pipeline_tag in task.pipeline_tags]


def compatible_text_generation_frameworks(
    target_gpu: str,
    supported_frameworks: set[str],
    inspection: ModelInspection,
) -> list[str]:
    can_llamacpp = "llamacpp" in supported_frameworks
    vllm_like = [framework for framework in VLLM_LIKE_FRAMEWORKS if framework in supported_frameworks]
    can_transformers = "transformers" in supported_frameworks
    compatible: list[str] = []

    if can_llamacpp and inspection.has_gguf:
        compatible.append("llamacpp")
    if inspection.has_vllm_weights:
        compatible.extend(vllm_like)
        if can_transformers:
            compatible.append("transformers")
    if not compatible:
        raise ValueError(f"No compatible LLM weights/framework combination is available for {target_gpu}")
    return compatible


def choose_text_generation_framework(target_gpu: str, supported_frameworks: set[str], inspection: ModelInspection) -> str:
    return compatible_text_generation_frameworks(target_gpu, supported_frameworks, inspection)[0]


def compatible_frameworks_for_task(
    task_type: str,
    target_gpu: str,
    supported_frameworks: set[str],
    inspection: ModelInspection,
) -> list[str]:
    if task_type in {"text-generation", "visual-multi-modal", "reinforcement_learning"}:
        return compatible_text_generation_frameworks(target_gpu, supported_frameworks, inspection)

    if task_type == "asr":
        compatible: list[str] = []
        if "sherpa-onnx" in supported_frameworks and inspection.has_onnx_weights:
            compatible.append("sherpa-onnx")
        if inspection.has_standard_weights:
            compatible.extend(framework for framework in ("transformers", "funasr") if framework in supported_frameworks)
        if compatible:
            return compatible
        raise ValueError(f"No compatible ASR framework found for {target_gpu}")

    if task_type == "feature_emb":
        if inspection.has_standard_weights:
            compatible = [framework for framework in ("sentence-transformers", "transformers") if framework in supported_frameworks]
            if compatible:
                return compatible
        raise ValueError(f"No compatible embedding framework found for {target_gpu}")

    if task_type in {"question_answering", "vision_classification", "text_classification"}:
        if "transformers" in supported_frameworks and inspection.has_standard_weights:
            return ["transformers"]
        raise ValueError(f"No compatible transformers template found for {task_type} on {target_gpu}")

    if task_type == "text-to-image-generation":
        if "diffusers" in supported_frameworks and inspection.has_standard_weights:
            return ["diffusers"]
        raise ValueError(f"No compatible diffusers template found for {target_gpu}")

    raise ValueError(f"Unsupported task type for auto framework selection: {task_type}")


def choose_framework_for_task(task_type: str, target_gpu: str, supported_frameworks: set[str], inspection: ModelInspection) -> str:
    return compatible_frameworks_for_task(task_type, target_gpu, supported_frameworks, inspection)[0]