import copy
import os
import re


METADATA_KEY = "metadata"


_IMAGE_PROMPT_LABELS = {
    "": "Text/new generation",
    "T": "Text/new generation",
    "S": "Start image",
    "E": "End image",
    "V": "Continue from source video",
    "L": "Continue from last generated video",
}


def get_model_family(base_model_type, model_def, model_types_handlers, families_infos, for_ui=False):
    if for_ui:
        model_family = model_def.get("group", None)
        if model_family is not None and model_family in families_infos:
            return model_family
    handler = model_types_handlers.get(base_model_type, None)
    if handler is None:
        return "unknown"
    return handler.query_model_family()


def _model_def_finetune(model_def):
    return os.path.basename(os.path.dirname(str(model_def.get("path", "") or ""))).casefold() == "finetunes"


def _choice_values(choice_def):
    if not isinstance(choice_def, dict):
        return []
    values = []
    for choice in choice_def.get("choices", []) or []:
        if isinstance(choice, (list, tuple)):
            values.append(str(choice[1] if len(choice) > 1 else choice[0] if len(choice) > 0 else "") or "")
        else:
            values.append(str(choice or ""))
    values.extend(str(value or "") for value in choice_def.get("selection", []) or [])
    default = choice_def.get("default", None)
    if default is not None:
        values.append(str(default or ""))
    return values


def _choice_values_contain(choice_def, letters):
    return any(any(letter in value for letter in letters) for value in _choice_values(choice_def))


def _normalized_choices(choice_def):
    if not isinstance(choice_def, dict):
        return []
    labels = choice_def.get("labels", {}) or {}
    choices = []
    seen = set()

    def add_choice(label, value):
        value = str(value or "")
        if value in seen:
            return
        seen.add(value)
        choices.append({"label": str(label or value or "None"), "value": value})

    for choice in choice_def.get("choices", []) or []:
        if isinstance(choice, (list, tuple)):
            label = choice[0] if len(choice) > 0 else ""
            value = choice[1] if len(choice) > 1 else label
            add_choice(label, value)
        else:
            add_choice(labels.get(choice, choice), choice)
    for key, label in labels.items():
        add_choice(label, key.replace("V", "").replace("P", ""))
    for value in choice_def.get("selection", []) or []:
        add_choice(labels.get(value, value), value)
    return choices


def normalize_choice_def(choice_def):
    if not isinstance(choice_def, dict):
        return None
    keys = ("default", "label", "letters_filter", "visible", "show_label", "scale", "trigger", "type")
    normalized = {key: choice_def[key] for key in keys if key in choice_def}
    normalized["choices"] = _normalized_choices(choice_def)
    return normalized


def infer_main_outputs(model_def):
    if model_def.get("audio_only", False):
        return ["audio"]
    if model_def.get("image_outputs", False):
        return ["image"]
    if model_def.get("v2i_switch_supported", False) or model_def.get("inpaint_support", False):
        return ["image", "video"]
    return ["video"]


def infer_outputs(model_def):
    outputs = infer_main_outputs(model_def)
    if model_def.get("returns_audio", False) and "audio" not in outputs:
        outputs.append("audio")
    return outputs


def infer_inputs(model_def):
    inputs = ["text"]
    image_prompt_types_allowed = str(model_def.get("image_prompt_types_allowed", "") or "")
    image_refs = model_def.get("image_ref_choices", None)
    alt_guide_refs = model_def.get("guide_custom_choices", None)
    guide_preprocessing = model_def.get("guide_preprocessing", None)
    custom_video_selection = model_def.get("custom_video_selection", None)
    image_outputs = bool(model_def.get("image_outputs", False))

    if model_def.get("any_audio_prompt", False):
        inputs.append("audio")
    if "S" in image_prompt_types_allowed or "E" in image_prompt_types_allowed or model_def.get("end_frames_always_enabled", False) or model_def.get("inpaint_support", False) or _choice_values_contain(image_refs, "IKF") or _choice_values_contain(alt_guide_refs, "IKF"):
        inputs.append("image")
    if image_outputs and (_choice_values_contain(guide_preprocessing, "V") or _choice_values_contain(alt_guide_refs, "V")) and "image" not in inputs:
        inputs.append("image")
    if "V" in image_prompt_types_allowed or "L" in image_prompt_types_allowed or (not image_outputs and (_choice_values_contain(guide_preprocessing, "V") or _choice_values_contain(alt_guide_refs, "V") or _choice_values_contain(custom_video_selection, "V"))):
        inputs.append("video")
    return list(dict.fromkeys(inputs))


def infer_media_inputs(model_def):
    image_prompt_types_allowed = str(model_def.get("image_prompt_types_allowed", "") or "")
    image_refs = model_def.get("image_ref_choices", None)
    guide_preprocessing = model_def.get("guide_preprocessing", None)
    guide_custom_choices = model_def.get("guide_custom_choices", None)
    custom_video_selection = model_def.get("custom_video_selection", None)
    has_reference = _choice_values_contain(image_refs, "I") or _choice_values_contain(guide_custom_choices, "I")
    single_reference = bool(model_def.get("one_image_ref_needed", False) or model_def.get("one_image_ref_only", False))
    has_control = _choice_values_contain(guide_preprocessing, "V") or _choice_values_contain(guide_custom_choices, "V") or _choice_values_contain(custom_video_selection, "V")
    image_outputs = bool(model_def.get("image_outputs", False))
    return {
        "image": {
            "start": "S" in image_prompt_types_allowed,
            "end": "E" in image_prompt_types_allowed or bool(model_def.get("end_frames_always_enabled", False)),
            "reference": has_reference,
            "single_reference": has_reference and single_reference,
            "multiple_references": has_reference and not single_reference,
            "background": _choice_values_contain(image_refs, "K") or _choice_values_contain(guide_custom_choices, "K"),
            "injected_frames": _choice_values_contain(image_refs, "F") or _choice_values_contain(guide_custom_choices, "F"),
            "control": image_outputs and has_control,
            "mask": image_outputs and (_choice_values_contain(model_def.get("mask_preprocessing", None), "A") or bool(model_def.get("inpaint_support", False))),
        },
        "video": {
            "continue": "V" in image_prompt_types_allowed,
            "last": "L" in image_prompt_types_allowed,
            "control": (not image_outputs) and has_control,
            "mask": (not image_outputs) and _choice_values_contain(model_def.get("mask_preprocessing", None), "A"),
        },
        "audio": {
            "prompt": bool(model_def.get("any_audio_prompt", False)),
            "output": bool(model_def.get("audio_only", False) or model_def.get("returns_audio", False)),
        },
    }


def infer_capabilities(model_def, main_outputs, outputs, inputs, media_inputs):
    image_inputs = media_inputs["image"]
    video_inputs = media_inputs["video"]
    audio_inputs = media_inputs["audio"]
    return {
        "text_to_video": "video" in main_outputs and "text" in inputs,
        "image_to_video": "video" in main_outputs and image_inputs["start"],
        "video_to_video": "video" in main_outputs and (video_inputs["continue"] or video_inputs["control"]),
        "text_to_image": "image" in main_outputs and "text" in inputs,
        "image_to_image": "image" in main_outputs and (image_inputs["start"] or image_inputs["reference"] or image_inputs["control"]),
        "text_to_audio": "audio" in main_outputs and "text" in inputs,
        "audio_to_audio": "audio" in main_outputs and audio_inputs["prompt"],
        "audio_to_video": "video" in main_outputs and audio_inputs["prompt"],
        "audio_output": "audio" in outputs,
        "inpainting": bool(model_def.get("inpaint_support", False) or image_inputs["mask"] or video_inputs["mask"]),
        "outpainting": bool(model_def.get("video_guide_outpainting", False)),
        "reference_images": image_inputs["reference"],
        "background_image": image_inputs["background"],
        "injected_frames": image_inputs["injected_frames"],
        "control_image": image_inputs["control"],
        "control_video": video_inputs["control"],
        "video_continuation": video_inputs["continue"],
        "sliding_window": bool(model_def.get("sliding_window", False)),
        "lora": not bool(model_def.get("no_lora", False)),
    }


def infer_setting_values(model_def):
    image_prompt_types_allowed = str(model_def.get("image_prompt_types_allowed", "") or "")
    image_prompt_choices = [{"label": _IMAGE_PROMPT_LABELS.get("", ""), "value": ""}]
    for letter in image_prompt_types_allowed:
        if letter == "T":
            continue
        image_prompt_choices.append({"label": _IMAGE_PROMPT_LABELS.get(letter, letter), "value": letter})
    return {
        "image_prompt_type": {
            "allowed": image_prompt_types_allowed,
            "choices": image_prompt_choices,
        },
        "video_prompt_type": {
            "guide_preprocessing": normalize_choice_def(model_def.get("guide_preprocessing", None)),
            "mask_preprocessing": normalize_choice_def(model_def.get("mask_preprocessing", None)),
            "guide_custom_choices": normalize_choice_def(model_def.get("guide_custom_choices", None)),
            "image_ref_choices": normalize_choice_def(model_def.get("image_ref_choices", None)),
            "custom_video_selection": normalize_choice_def(model_def.get("custom_video_selection", None)),
            "forced": str(model_def.get("set_video_prompt_type", "") or ""),
        },
        "audio_prompt_type": {
            "sources": normalize_choice_def(model_def.get("audio_prompt_type_sources", None)),
            "custom_option": model_def.get("audio_prompt_type_custom_option", None),
        },
        "model_mode": normalize_choice_def(model_def.get("model_modes", None)),
        "sample_solver": normalize_choice_def({"choices": model_def.get("sample_solvers", [])}) if model_def.get("sample_solvers", None) is not None else None,
        "prompt_enhancer": normalize_choice_def(model_def.get("prompt_enhancer_def", None)),
    }


def store_metadata(model_type, model_def, model_types_handlers, families_infos):
    base_model_type = model_def.get("architecture", None) or model_type
    family = get_model_family(base_model_type, model_def, model_types_handlers, families_infos, for_ui=True)
    family_label = families_infos.get(family, families_infos.get("unknown", (100, "Unknown")))[1]
    main_outputs = infer_main_outputs(model_def)
    outputs = infer_outputs(model_def)
    inputs = infer_inputs(model_def)
    media_inputs = infer_media_inputs(model_def)
    model_def[METADATA_KEY] = {
        "model_type": model_type,
        "family": family,
        "family_label": family_label,
        "base_model_type": base_model_type,
        "finetune": _model_def_finetune(model_def),
        "main_output": main_outputs,
        "outputs": outputs,
        "inputs": inputs,
        "media_inputs": media_inputs,
        "capabilities": infer_capabilities(model_def, main_outputs, outputs, inputs, media_inputs),
        "setting_values": infer_setting_values(model_def),
    }
    return model_def


def _normalize_filter_values(value, *, split_string=True):
    if value is None:
        return None
    if isinstance(value, (list, tuple, set)):
        values = [str(one).strip() for one in value if str(one).strip()]
    elif split_string and isinstance(value, str):
        values = [one.strip() for one in re.split(r"[,|/]", value) if one.strip()]
    else:
        values = [str(value).strip()] if str(value).strip() else []
    return values or None


def _normalize_bool_filter(value):
    if value is None or isinstance(value, str) and len(value.strip()) == 0:
        return None
    if isinstance(value, bool):
        return value
    text = str(value).strip().casefold()
    if text in ("1", "true", "yes", "y", "on"):
        return True
    if text in ("0", "false", "no", "n", "off"):
        return False
    raise ValueError(f"Invalid finetune filter value: {value}")


def _metadata_matches_filter(value, expected):
    expected = _normalize_filter_values(expected)
    if expected is None:
        return True
    if isinstance(value, (list, tuple, set)):
        actual = {str(one).casefold() for one in value}
        return any(one.casefold() in actual for one in expected)
    actual = str(value or "").casefold()
    return any(actual == one.casefold() for one in expected)


def list_model_defs(models_def, *, family=None, base_model_type=None, finetune=None, model_type=None, main_output=None, inputs=None):
    finetune_filter = _normalize_bool_filter(finetune)
    requested_outputs = _normalize_filter_values(main_output)
    requested_inputs = _normalize_filter_values(inputs)
    requested_model_types = _normalize_filter_values(model_type, split_string=False)
    records = []
    for one_model_type, model_def in models_def.items():
        metadata = model_def.get(METADATA_KEY, {})
        if requested_model_types is not None and one_model_type not in requested_model_types:
            continue
        if not _metadata_matches_filter(metadata.get("family"), family) and not _metadata_matches_filter(metadata.get("family_label"), family):
            continue
        if not _metadata_matches_filter(metadata.get("base_model_type"), base_model_type):
            continue
        if finetune_filter is not None and bool(metadata.get("finetune", False)) != finetune_filter:
            continue
        if requested_outputs is not None and not _metadata_matches_filter(metadata.get("main_output", []), requested_outputs):
            continue
        if requested_inputs is not None and not _metadata_matches_filter(metadata.get("inputs", []), requested_inputs):
            continue
        record = copy.deepcopy(model_def)
        record["model_type"] = one_model_type
        records.append(record)
    return records