#!/usr/bin/env python3
"""Find model-like files hosted on GitHub.

This script ONLY talks to GitHub (api.github.com) and is designed for
restricted enterprise networks where only GitHub is reachable.

It searches repositories by keyword, then inspects recent releases to find
assets that look like model files (e.g., .gguf, .safetensors, .bin, .zip).

Usage:
    python oss_llm/find_models_on_github.py --query "qwen2.5 gguf" --limit 10

Optional:
  - Set GITHUB_TOKEN to increase API rate limits.

Notes:
  - GitHub rarely hosts very large model weights due to size constraints.
  - Prefer smaller models and quantized artifacts.
"""

import argparse
import os
import sys
import textwrap
import urllib.parse
import json
import subprocess

MODEL_EXTENSIONS = (
    ".gguf",
    ".safetensors",
    ".bin",
    ".pt",
    ".pth",
    ".onnx",
    ".zip",
    ".tar.gz",
    ".tgz",
)


def _http_get_json(url: str) -> object:
    token = os.environ.get("GITHUB_TOKEN") or os.environ.get("GH_TOKEN")

    cmd = [
        "curl",
        "-L",
        "--fail",
        "--silent",
        "--show-error",
        "-H",
        "Accept: application/vnd.github+json",
    ]
    if token:
        cmd += ["-H", f"Authorization: Bearer {token}"]
    cmd.append(url)

    proc = subprocess.run(cmd, capture_output=True, text=True)
    if proc.returncode != 0:
        raise RuntimeError(proc.stderr.strip() or f"curl failed with exit code {proc.returncode}")
    return json.loads(proc.stdout)


def _looks_like_model_asset(name: str) -> bool:
    lower = name.lower()
    return any(lower.endswith(ext) for ext in MODEL_EXTENSIONS)


def _safe_int(value: object) -> int | None:
    try:
        return int(value)  # type: ignore[arg-type]
    except Exception:
        return None


def _scan_repo_tree_for_models(full_name: str, default_branch: str) -> list[str]:
    """Return matching file paths from the repo tree.

    This is useful when model files are stored in-repo (often via Git LFS),
    and therefore don't show up as release assets.
    """

    # The Trees API can be large; callers should keep repo counts low.
    url = f"https://api.github.com/repos/{full_name}/git/trees/{urllib.parse.quote(default_branch)}?recursive=1"
    data = _http_get_json(url)
    if not isinstance(data, dict):
        return []
    tree = data.get("tree")
    if not isinstance(tree, list):
        return []

    hits: list[str] = []
    for node in tree:
        if not isinstance(node, dict):
            continue
        path = node.get("path")
        ntype = node.get("type")
        if ntype != "blob" or not isinstance(path, str):
            continue
        if _looks_like_model_asset(path):
            hits.append(path)
    return hits


def _human_size(num_bytes: int) -> str:
    n = float(num_bytes)
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if n < 1024.0:
            return f"{n:.1f}{unit}"
        n /= 1024.0
    return f"{n:.1f}PB"


def main() -> int:
    p = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description="Search GitHub releases for model-like assets.",
        epilog=textwrap.dedent(
            """
            Examples:
                            python oss_llm/find_models_on_github.py --query "llama gguf" --limit 20
                            python oss_llm/find_models_on_github.py --query "qwen 2.5 gguf" --limit 10

            Tips:
              - Add qualifiers to narrow results, e.g. `language:python`, `in:name`, `topic:llama-cpp`.
              - Set GITHUB_TOKEN to avoid low unauthenticated rate limits.
            """
        ),
    )
    p.add_argument("--query", required=True, help="GitHub search query")
    p.add_argument("--limit", type=int, default=10, help="Max repositories to inspect")
    p.add_argument(
        "--per-page",
        type=int,
        default=10,
        help="Repos per page from search API (max 100)",
    )
    p.add_argument(
        "--max-releases",
        type=int,
        default=5,
        help="Max recent releases per repo to inspect",
    )
    p.add_argument(
        "--scan-tree",
        action="store_true",
        help="Also scan repository file trees for model-like files (slower, more API calls)",
    )
    args = p.parse_args()

    query = args.query.strip()
    if not query:
        print("--query must be non-empty", file=sys.stderr)
        return 2

    per_page = max(1, min(100, args.per_page))
    limit = max(1, args.limit)

    search_q = urllib.parse.quote(query)
    search_url = f"https://api.github.com/search/repositories?q={search_q}&per_page={per_page}"

    try:
        search = _http_get_json(search_url)
    except Exception as e:
        print(f"GitHub search failed: {e}", file=sys.stderr)
        return 1

    items = search.get("items") if isinstance(search, dict) else None
    if not items:
        print("No repositories found.")
        return 0

    inspected = 0
    found_any = False

    for repo in items:
        if inspected >= limit:
            break

        full_name = repo.get("full_name")
        html_url = repo.get("html_url")
        default_branch = repo.get("default_branch")
        if not full_name or not html_url:
            continue

        inspected += 1
        print(f"\n== {full_name} ==")
        print(html_url)

        releases_url = f"https://api.github.com/repos/{full_name}/releases?per_page={args.max_releases}"
        try:
            releases = _http_get_json(releases_url)
        except Exception as e:
            print(f"  releases: error: {e}")
            continue

        if not isinstance(releases, list) or len(releases) == 0:
            print("  releases: none")
            releases = []

        repo_hit = False
        for rel in releases[: args.max_releases]:
            tag = rel.get("tag_name") or "(no tag)"
            name = rel.get("name") or "(no name)"
            assets = rel.get("assets") or []
            model_assets = [a for a in assets if _looks_like_model_asset(a.get("name", ""))]

            if not model_assets:
                continue

            repo_hit = True
            found_any = True
            print(f"  release: {tag} — {name}")
            for a in model_assets:
                aname = a.get("name") or "(no name)"
                size = a.get("size")
                url = a.get("browser_download_url") or ""
                size_str = _human_size(int(size)) if isinstance(size, int) else "?"
                print(f"    - {aname} ({size_str})")
                if url:
                    print(f"      {url}")

        if releases and not repo_hit:
            print("  releases: present, but no model-like assets found")

        if args.scan_tree and isinstance(default_branch, str) and default_branch:
            try:
                paths = _scan_repo_tree_for_models(full_name, default_branch)
            except Exception as e:
                print(f"  repo files: error: {e}")
                continue

            if paths:
                found_any = True
                print(f"  repo files: found {len(paths)} model-like paths")
                for pth in paths[:30]:
                    print(f"    - {pth}")
                if len(paths) > 30:
                    print("    - (more omitted)")
            else:
                print("  repo files: no model-like paths found")

    if not found_any:
        print("\nNo model-like release assets found in inspected repositories.")
        print("Try refining your query (add 'gguf', model family name, or 'quant').")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())
