import importlib.util
import transformers
import diffusers
from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae, errors
from modules.logger import log
from pipelines import generic


def _import_from_file(module_name, file_path):
    spec = importlib.util.spec_from_file_location(module_name, file_path)
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


def init_transformer_component(repo_id, diffusers_load_config, adapter_cls):
    """Load (transformer, llm_adapter_or_none).

    If the UNET dropdown points at a valid safetensors, route through
    :mod:`pipelines.native_transformer` with :data:`pipelines.anima.ANIMA_SPEC`,
    which extracts any bundled ``llm_adapter`` weights inline with the
    transformer. Otherwise fall back to :func:`generic.load_transformer` and
    return ``None`` for the adapter so the caller loads it from the base repo.
    """
    from pipelines import native_transformer
    local_file = native_transformer.resolve_path()
    if local_file is not None:
        from pipelines.anima import ANIMA_SPEC
        try:
            transformer, siblings = native_transformer.load(
                local_file, repo_id, ANIMA_SPEC, diffusers_load_config,
                sibling_classes={'llm_adapter': adapter_cls},
            )
            return transformer, siblings.get('llm_adapter')
        except Exception as e:
            log.error(f'Load model: type=Anima custom transformer="{local_file}": {e}')
            errors.display(e, 'Load')
            return None, None
    transformer = generic.load_transformer(
        repo_id,
        cls_name=diffusers.CosmosTransformer3DModel,
        load_config=diffusers_load_config,
        subfolder="transformer"
    )
    return transformer, None


def load_anima(checkpoint_info, diffusers_load_config=None):
    if diffusers_load_config is None:
        diffusers_load_config = {}
    repo_id = sd_models.path_to_repo(checkpoint_info)
    sd_models.hf_auth_check(checkpoint_info)

    load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False)
    load_args.pop('cache_dir', None)
    log.debug(f'Load model: type=Anima repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}')

    if repo_id is None or repo_id.lower() == 'none':
        return None

    # load-or-download custom pipeline modules from repo
    """
    import os
    import sys
    import huggingface_hub as hf

    if os.path.exists(os.path.join(repo_id, 'pipeline.py')):
        pipeline_file = os.path.join(repo_id, 'pipeline.py')
    else:
        try:
            if os.path.exists(repo_id):
                from pipelines.generic_map import transformers_map
                custom_id = transformers_map.get('AnimaTextToImagePipeline', repo_id)
            else:
                custom_id = repo_id
            pipeline_file = hf.hf_hub_download(repo_id=custom_id, filename='pipeline.py', cache_dir=shared.opts.hfcache_dir)
        except Exception as e:
            log.error(f'Load model: type=Anima failed to download custom modules: {e}')
            return None
    if os.path.exists(os.path.join(repo_id, 'llm_adapter/modeling_llm_adapter.py')):
        adapter_file = os.path.join(repo_id, 'llm_adapter/modeling_llm_adapter.py')
    else:
        try:
            if os.path.exists(repo_id):
                from pipelines.generic_map import transformers_map
                custom_id = transformers_map.get('AnimaTextToImagePipeline', repo_id)
            else:
                custom_id = repo_id
            adapter_file = hf.hf_hub_download(repo_id=custom_id, filename='llm_adapter/modeling_llm_adapter.py', cache_dir=shared.opts.hfcache_dir)
        except Exception as e:
            log.error(f'Load model: type=Anima failed to download custom modules: {e}')
            return None

    # dynamically import custom classes and register in sys.modules so Diffusers' from_pretrained can resolve them via trust_remote_code
    adapter_mod = _import_from_file('modeling_llm_adapter', adapter_file)
    sys.modules['modeling_llm_adapter'] = adapter_mod
    pipeline_mod = _import_from_file('pipeline', pipeline_file)
    sys.modules['pipeline'] = pipeline_mod
    AnimaTextToImagePipeline = pipeline_mod.AnimaTextToImagePipeline
    AnimaLLMAdapter = adapter_mod.AnimaLLMAdapter
    """

    import sys
    from pipelines.anima import modeling_llm_adapter
    sys.modules['modeling_llm_adapter'] = modeling_llm_adapter
    from pipelines.anima.pipeline import AnimaTextToImagePipeline
    from pipelines.anima.anima_image import build_anima_pipeline_classes
    AnimaImageToImagePipeline, AnimaInpaintPipeline = build_anima_pipeline_classes(AnimaTextToImagePipeline)
    diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["anima"] = AnimaTextToImagePipeline
    diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["anima"] = AnimaImageToImagePipeline
    diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["anima"] = AnimaInpaintPipeline
    generic.set_pipeline('Anima', AnimaTextToImagePipeline)

    # UNET dropdown (shared.opts.sd_unet) may redirect the transformer to a
    # community file that bundles both the transformer and the llm_adapter.
    transformer, llm_adapter = init_transformer_component(repo_id, diffusers_load_config, modeling_llm_adapter.AnimaLLMAdapter)
    if transformer is None:
        return None
    text_encoder = generic.load_text_encoder(
        repo_id,
        cls_name=transformers.Qwen3Model,
        load_config=diffusers_load_config,
        subfolder="text_encoder"
    )

    if llm_adapter is None:
        shared.state.begin('Load adapter')
        try:
            llm_adapter = modeling_llm_adapter.AnimaLLMAdapter.from_pretrained(
                repo_id,
                subfolder="llm_adapter",
                cache_dir=shared.opts.hfcache_dir,
                torch_dtype=devices.dtype,
            )
        except Exception as e:
            log.error(f'Load model: type=Anima adapter: {e}')
            return None
        finally:
            shared.state.end()

    # assemble pipeline
    pipe = AnimaTextToImagePipeline.from_pretrained(
        repo_id,
        transformer=transformer,
        text_encoder=text_encoder,
        llm_adapter=llm_adapter,
        cache_dir=shared.opts.diffusers_dir,
        trust_remote_code=True,
        **load_args,
    )

    generic.load_vae_override(pipe, diffusers_load_config)

    del text_encoder
    del transformer
    del llm_adapter

    sd_hijack_te.init_hijack(pipe)
    sd_hijack_vae.init_hijack(pipe)

    devices.torch_gc()
    return pipe