from .imports import *
# ---------------------------------------------------------------------------
# Resolvers — each is independent; failure returns {} not an exception
# ---------------------------------------------------------------------------

def resolve_local_config(directory: str, hub_id: str) -> dict:
    path = os.path.join(directory, "config.json")
    if not os.path.isfile(path):
        return {}
    try:
        with open(path, "r") as f:
            cfg = json.load(f)
    except (OSError, json.JSONDecodeError):
        return {}
    return {
        "architectures":           cfg.get("architectures"),
        "model_type":              cfg.get("model_type"),
        "max_position_embeddings": cfg.get("max_position_embeddings"),
        "sliding_window":          cfg.get("sliding_window"),
        "rope_scaling":            cfg.get("rope_scaling"),
        "torch_dtype":             cfg.get("torch_dtype"),
        "vocab_size":              cfg.get("vocab_size"),
    }


def resolve_local_tokenizer(directory: str, hub_id: str) -> dict:
    path = os.path.join(directory, "tokenizer_config.json")
    if not os.path.isfile(path):
        return {}
    try:
        with open(path, "r") as f:
            tk = json.load(f)
    except (OSError, json.JSONDecodeError):
        return {}
    raw = tk.get("model_max_length")
    if isinstance(raw, (int, float)) and 0 < raw < TOKENIZER_SENTINEL_THRESHOLD:
        return {"tokenizer_model_max_length": int(raw)}
    return {}


def resolve_hub_model_info(directory: str, hub_id: str, api: HfApi) -> dict:
    try:
        info = api.model_info(hub_id, files_metadata=False)
    except HfHubHTTPError:
        return {}

    # safetensors metadata is None for GGUF/diffusers/etc. — guard.
    params = getattr(info.safetensors, "total", None) if info.safetensors else None
    auto_model = getattr(info.transformers_info, "auto_model", None) if info.transformers_info else None

    # card_data can be dict-like or a typed CardData; normalize.
    card = info.card_data
    def _card(key):
        if card is None:
            return None
        return card.get(key) if isinstance(card, dict) else getattr(card, key, None)

    languages = _card("language")
    if isinstance(languages, str):
        languages = [languages]

    return {
        "pipeline_tag":     info.pipeline_tag,
        "library_name":     info.library_name,
        "auto_model_class": auto_model,
        "parameter_count":  params,
        "license":          _card("license"),
        "gated":            bool(info.gated) if info.gated is not None else None,
        "languages":        languages,
        "tags":             info.tags,
    }


# ---------------------------------------------------------------------------
# Registry of resolvers, ordered by trust
# ---------------------------------------------------------------------------

ResolverFn = Callable[[str, str], dict]


def build_resolver_chain(
    *,
    api: Optional[HfApi] = None,
    use_hub: bool = True,
) -> List[Tuple[str, ResolverFn]]:
    """Explicit chain — caller can drop resolvers, reorder, or insert.
    No module-level state; the HfApi instance is passed in."""
    chain: List[Tuple[str, ResolverFn]] = [
        ("local_config",    resolve_local_config),
        ("local_tokenizer", resolve_local_tokenizer),
    ]
    if use_hub:
        hub_api = api or HfApi()
        chain.append(
            ("hub_model_info", lambda d, h: resolve_hub_model_info(d, h, hub_api))
        )
    return chain


def enrich(
    directory: str,
    hub_id: str,
    chain: List[Tuple[str, ResolverFn]],
) -> Tuple[ModelMetadata, Dict[str, str]]:
    """Run the chain. Returns the typed record + a {field: source} provenance map."""
    valid_fields = {f.name for f in dataclass_fields(ModelMetadata)}
    merged: Dict[str, Any] = {"hub_id": hub_id}
    sources: Dict[str, str] = {"hub_id": "input"}

    for source_name, fn in chain:
        partial = fn(directory, hub_id)
        for field_name, value in partial.items():
            if value is None or field_name not in valid_fields:
                continue
            if merged.get(field_name) is None:
                merged[field_name] = value
                sources[field_name] = source_name

    return ModelMetadata(**merged), sources

def discover_model(save_json=True,verbose=True,use_hub=True):
    discovered={}
    model_dirs = get_model_dirs()
    chain = build_resolver_chain(use_hub=use_hub)
    for directory in model_dirs:
        file_parts = get_file_parts(directory)
        folder = directory[len(MODELS_HOME):]
        
        name = file_parts.get("basename")
        parent_dirname = file_parts.get("parent_dirname")
        hub_id = directory[len(parent_dirname)+1:]
        meta, meta_sources = enrich(directory, hub_id, chain)
        discovered[name] = {"name":name,"hub_id":hub_id,"folder":folder}
        discovered[name].update({
            "pipeline_tag":            meta.pipeline_tag,
            "library_name":            meta.library_name,
            "auto_model_class":        meta.auto_model_class,
            "architectures":           meta.architectures,
            "model_type":              meta.model_type,
            "max_position_embeddings": meta.max_position_embeddings,
            "model_max_length":        meta.tokenizer_model_max_length or meta.max_position_embeddings or DEFAULT_MAX_TOKENS,
            "parameter_count":         meta.parameter_count,
            "license":                 meta.license,
            "gated":                   meta.gated,
            "languages":               meta.languages,
        })
        input(discovered)
        if verbose:
            print(f"[enrich] {hub_id}: {meta_sources}")
    if save_json:
        safe_dump_to_file(data=discovered,file_path=MODELS_DICT_PATH)
        
    return discovered
def discover_models(save_json=True, verbose=True, use_hub=True):
    """Walk the model tree and enrich each dir. Descriptive only.
    Writes a REPORT, never the overlay — acquire() owns MODELS_DICT_PATH."""
    discovered = {}
    chain = build_resolver_chain(use_hub=use_hub)

    for directory in get_model_dirs():
        file_parts = get_file_parts(directory)
        name = file_parts.get("basename")
        hub_id = directory[len(file_parts.get("parent_dirname")) + 1:]

        meta, meta_sources = enrich(directory, hub_id, chain)
        discovered[name] = meta.to_dict()        # the typed record, whole
        if verbose:
            print(f"[enrich] {hub_id}: {meta_sources}")

    if save_json:
        safe_dump_to_file(data=discovered, file_path=MODELS_DISCOVERY_PATH)  # NOT MODELS_DICT_PATH
    return discovered
