#!/usr/bin/env python
"""
Inspect Litellm model metadata and support flags.

This mirrors litellm.utils._supports_factory internals to show everything it
discovers (resolved provider, model_info, provider_info fallbacks, etc.).
"""

from __future__ import annotations

import argparse
import json
import sys
from typing import Any

import litellm
from litellm.utils import (
    _get_model_info_helper,
    _supports_provider_info_factory,
    supports_vision,
)
from openhands.sdk import LLM


def dump_supports_factory(
    model: str, custom_llm_provider: str | None, key: str
) -> bool:
    """Verbose version of litellm.utils._supports_factory."""
    try:
        resolved_model, resolved_provider, _, _ = litellm.get_llm_provider(
            model=model, custom_llm_provider=custom_llm_provider
        )
        print(f"Resolved model/provider: {resolved_model} / {resolved_provider}")

        model_info: dict[str, Any] = _get_model_info_helper(
            model=resolved_model, custom_llm_provider=resolved_provider
        )
        print("model_info:")
        print(json.dumps(model_info, indent=2, ensure_ascii=False))

        if model_info.get(key, False) is True:
            print(f"model_info[{key}] == True -> supported")
            return True
        if model_info.get(key) is None:
            provider_support = _supports_provider_info_factory(
                resolved_model, resolved_provider, key
            )
            print(f"provider_support[{key}] fallback: {provider_support}")
            if provider_support is not None:
                return provider_support

        print(f"{key} not supported per model/provider info")
        return False
    except Exception as e:  # pragma: no cover - debug utility
        print(
            f"[warn] Error checking {key} support for model={model}, provider={custom_llm_provider}: {e}"
        )
        provider_support = _supports_provider_info_factory(
            model, custom_llm_provider, key
        )
        print(f"provider_support[{key}] fallback after exception: {provider_support}")
        if provider_support is not None:
            return provider_support
        return False


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--model", default="litellm_proxy/gpt-5-mini", help="Model name to inspect."
    )
    parser.add_argument(
        "--provider", default=None, help="Custom provider name (optional)."
    )
    parser.add_argument(
        "--key",
        default="supports_function_calling",
        help="Support key to check (e.g., supports_function_calling, supports_vision, native_tool_calling).",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()

    print(f"=== Checking model: {args.model} (provider={args.provider}) ===")
    supported = dump_supports_factory(args.model, args.provider, args.key)
    print(f"\nResult for {args.key}: {supported}")

    # Extra: vision flag via litellm.utils.supports_vision
    vis = supports_vision(args.model)
    print(f"supports_vision (litellm.utils.supports_vision): {vis}")

    # Extra: top-level model listing info if available
    try:
        _, provider, _, _ = litellm.get_llm_provider(
            model=args.model, custom_llm_provider=args.provider
        )
        print(f"Resolved provider: {provider}")
    except Exception as exc:  # pragma: no cover - debug output
        print(f"[warn] get_llm_provider failed: {exc}")

    llm = LLM(model=args.model)
    llm._init_model_info_and_caps()
    print(llm.model_info)

    return 0


if __name__ == "__main__":
    sys.exit(main())