#!/usr/bin/env python3
"""
Bioconda package crawler with tiered export (T0 / T1 / T2).

Data sources:
- https://conda.anaconda.org/bioconda/<subdir>/repodata.json
- https://api.anaconda.org/package/bioconda/<package_name>
"""

from __future__ import annotations

import argparse
import json
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple

import requests
from packaging.version import InvalidVersion, Version


REPO_BASE = "https://conda.anaconda.org/bioconda"
ANACONDA_PACKAGE_API = "https://api.anaconda.org/package/bioconda/{package_name}"
ANACONDA_CHANNEL_PACKAGES_API = "https://api.anaconda.org/packages/bioconda"


T1_DOMAIN_KEYWORDS = {
    "single_cell": [
        "single-cell",
        "single cell",
        "scrna",
        "scanpy",
        "seurat",
        "cellranger",
        "scvelo",
        "scvi",
    ],
    "spatial_transcriptomics": [
        "spatial",
        "visium",
        "stereoseq",
        "spaceranger",
        "squidpy",
        "spatialde",
    ],
    "proteomics": [
        "proteomics",
        "proteome",
        "peptide",
        "mass spectrometry",
        "metaproteomics",
    ],
    "metabolomics": [
        "metabolomics",
        "metabolome",
        "metabolic profiling",
        "lipidomics",
        "lc-ms",
    ],
}


@dataclass
class PackageRecord:
    package_name: str
    latest_version: str
    build_number: int
    depends: List[str]
    subdir: str
    timestamp: int
    repodata_summary: str = ""
    downloads: int = -1


def _safe_version(raw_version: str) -> Version:
    try:
        return Version(raw_version)
    except InvalidVersion:
        return Version("0")


def fetch_repodata(subdir: str, timeout: int = 60) -> Dict:
    url = f"{REPO_BASE}/{subdir}/repodata.json"
    response = requests.get(url, timeout=timeout)
    response.raise_for_status()
    return response.json()


def iter_repodata_packages(payload: Dict) -> Iterable[Tuple[str, Dict]]:
    for section in ("packages", "packages.conda"):
        for filename, meta in payload.get(section, {}).items():
            yield filename, meta


def build_latest_index(subdirs: List[str]) -> Dict[str, PackageRecord]:
    latest: Dict[str, PackageRecord] = {}

    for subdir in subdirs:
        payload = fetch_repodata(subdir=subdir)
        for _filename, meta in iter_repodata_packages(payload):
            pkg_name = meta.get("name")
            version = meta.get("version", "")
            build_number = int(meta.get("build_number", 0))
            depends = meta.get("depends", []) or []
            timestamp = int(meta.get("timestamp", 0))
            repodata_summary = meta.get("summary", "") or ""
            if not pkg_name:
                continue

            candidate = PackageRecord(
                package_name=pkg_name,
                latest_version=version,
                build_number=build_number,
                depends=depends,
                subdir=subdir,
                timestamp=timestamp,
                repodata_summary=repodata_summary,
            )

            if pkg_name not in latest:
                latest[pkg_name] = candidate
                continue

            current = latest[pkg_name]
            candidate_key = (_safe_version(candidate.latest_version), candidate.build_number, candidate.timestamp)
            current_key = (_safe_version(current.latest_version), current.build_number, current.timestamp)
            if candidate_key > current_key:
                latest[pkg_name] = candidate

    return latest


def _normalize_owners(owners_field) -> List[str]:
    if not owners_field:
        return []
    owners: List[str] = []
    for item in owners_field:
        if isinstance(item, str):
            owners.append(item)
        elif isinstance(item, dict):
            login = item.get("login") or item.get("name")
            if login:
                owners.append(login)
    return sorted(set(owners))


def fetch_package_profile(package_name: str, timeout: int = 30) -> Dict:
    url = ANACONDA_PACKAGE_API.format(package_name=package_name)
    response = requests.get(url, timeout=timeout)
    if response.status_code == 404:
        return {}
    response.raise_for_status()
    return response.json()


def extract_download_count(profile: Dict[str, Any]) -> int:
    for key in ("ndownloads", "downloads", "download_count"):
        value = profile.get(key)
        if isinstance(value, int):
            return value
        if isinstance(value, str) and value.isdigit():
            return int(value)
    return -1


def _match_any_keyword(text: str, keywords: List[str]) -> bool:
    low = text.lower()
    return any(keyword in low for keyword in keywords)


def _decode_channel_packages_payload(payload: Any) -> Dict[str, int]:
    result: Dict[str, int] = {}
    if isinstance(payload, list):
        entries = payload
    elif isinstance(payload, dict):
        entries = payload.get("items") or payload.get("packages") or []
    else:
        entries = []

    for item in entries:
        if not isinstance(item, dict):
            continue
        pkg_name = item.get("name")
        if not pkg_name:
            continue
        result[pkg_name] = max(result.get(pkg_name, -1), extract_download_count(item))
    return result


def fetch_download_index_from_channel_api(timeout: int = 60, max_pages: int = 50) -> Dict[str, int]:
    download_index: Dict[str, int] = {}
    page = 1
    while True:
        params = {"page": page}
        response = requests.get(ANACONDA_CHANNEL_PACKAGES_API, params=params, timeout=timeout)
        response.raise_for_status()
        payload = response.json()
        current_page = _decode_channel_packages_payload(payload)

        if not current_page:
            if page == 1:
                current_page = _decode_channel_packages_payload(payload)
            if not current_page:
                break

        for pkg_name, downloads in current_page.items():
            download_index[pkg_name] = max(download_index.get(pkg_name, -1), downloads)

        if isinstance(payload, list):
            break

        total_pages = payload.get("total_pages")
        next_page = payload.get("next_page")
        if isinstance(total_pages, int) and page >= total_pages:
            break
        if next_page is None and not payload.get("items"):
            break
        page += 1
        if page > max_pages:
            break

    return download_index


def build_download_index_from_package_api(
    index: Dict[str, PackageRecord],
    profile_cache: Dict[str, Dict],
    timeout: int = 30,
    sleep_seconds: float = 0.05,
    max_packages: int = 0,
) -> Dict[str, int]:
    download_index: Dict[str, int] = {}
    package_names = sorted(index.keys())
    if max_packages > 0:
        package_names = package_names[:max_packages]

    for package_name in package_names:
        profile = profile_cache.get(package_name)
        if profile is None:
            profile = fetch_package_profile(package_name=package_name, timeout=timeout)
            profile_cache[package_name] = profile
            if sleep_seconds > 0:
                time.sleep(sleep_seconds)
        download_index[package_name] = extract_download_count(profile)
    return download_index


def resolve_download_index(
    index: Dict[str, PackageRecord],
    download_source: str,
    profile_cache: Dict[str, Dict],
    package_api_max: int,
) -> Dict[str, int]:
    print(f"max packages for package_api mode: {package_api_max}")
    if download_source == "channel_api":
        try:
            print("Fetching download index from channel API...")
            return fetch_download_index_from_channel_api()
            
        except Exception:
            print("Failed to fetch download index from channel API, falling back to package API...")
            return build_download_index_from_package_api(
                index=index,
                profile_cache=profile_cache,
                max_packages=package_api_max,
            )
    return build_download_index_from_package_api(
        index=index,
        profile_cache=profile_cache,
        max_packages=package_api_max,
    )


def _sorted_by_download(items: List[Tuple[str, PackageRecord]]) -> List[Tuple[str, PackageRecord]]:
    return sorted(
        items,
        key=lambda x: (x[1].downloads, _safe_version(x[1].latest_version), x[1].timestamp),
        reverse=True,
    )


def select_t0_records(
    index: Dict[str, PackageRecord],
    k_per_domain: int,
    k_overall: int,
) -> List[Tuple[str, PackageRecord, str]]:
    selected: List[Tuple[str, PackageRecord, str]] = []
    selected_names: Set[str] = set()

    for domain, keywords in T1_DOMAIN_KEYWORDS.items():
        domain_candidates: List[Tuple[str, PackageRecord]] = []
        for package_name, record in index.items():
            searchable = f"{package_name} {record.repodata_summary}"
            if _match_any_keyword(searchable, keywords):
                domain_candidates.append((package_name, record))

        ranked_domain = _sorted_by_download(domain_candidates)
        for package_name, record in ranked_domain[:k_per_domain]:
            if package_name in selected_names:
                continue
            selected.append((package_name, record, domain))
            selected_names.add(package_name)

    overall_candidates = _sorted_by_download(list(index.items()))
    for package_name, record in overall_candidates[:k_overall]:
        if package_name in selected_names:
            continue
        selected.append((package_name, record, "overall"))
        selected_names.add(package_name)

    return selected


def select_t1_records(index: Dict[str, PackageRecord], max_per_domain: int) -> List[Tuple[str, PackageRecord, str]]:
    selected: List[Tuple[str, PackageRecord, str]] = []
    used: Set[str] = set()

    for domain, keywords in T1_DOMAIN_KEYWORDS.items():
        counter = 0
        domain_candidates: List[Tuple[str, PackageRecord]] = []
        for package_name, record in index.items():
            if package_name in used:
                continue
            dep_text = " ".join(record.depends)
            searchable = f"{package_name} {record.repodata_summary} {dep_text}"
            if _match_any_keyword(searchable, keywords):
                domain_candidates.append((package_name, record))

        # Prefer higher-download tools first for better practical coverage.
        ranked_domain = _sorted_by_download(domain_candidates)
        for package_name, record in ranked_domain[:max_per_domain]:
            if package_name in used:
                continue
            selected.append((package_name, record, domain))
            used.add(package_name)
            counter += 1
            if counter >= max_per_domain:
                break
    return selected


def backfill_t1_records(
    index: Dict[str, PackageRecord],
    selected: List[Tuple[str, PackageRecord, str]],
    target_total: int,
) -> List[Tuple[str, PackageRecord, str]]:
    if target_total <= 0 or len(selected) >= target_total:
        return selected

    used = {pkg for pkg, _record, _domain in selected}
    ranked_overall = _sorted_by_download(list(index.items()))
    for package_name, record in ranked_overall:
        if package_name in used:
            continue
        selected.append((package_name, record, "t1_backfill_overall"))
        used.add(package_name)
        if len(selected) >= target_total:
            break
    return selected


def select_t2_records(index: Dict[str, PackageRecord], tools: List[str]) -> List[Tuple[str, PackageRecord]]:
    selected: List[Tuple[str, PackageRecord]] = []
    for tool in tools:
        normalized = tool.strip()
        if not normalized:
            continue
        lower_map = {name.lower(): name for name in index.keys()}
        hit = lower_map.get(normalized.lower())
        if hit:
            selected.append((normalized, index[hit]))
    return selected


def enrich_record(
    software_name: str,
    record: PackageRecord,
    tier: str,
    domain: Optional[str] = None,
    sleep_seconds: float = 0.2,
    profile_cache: Optional[Dict[str, Dict]] = None,
) -> Dict:
    profile_cache = profile_cache or {}
    profile = profile_cache.get(record.package_name)
    if profile is None:
        profile = fetch_package_profile(record.package_name)
        profile_cache[record.package_name] = profile
        if sleep_seconds > 0:
            time.sleep(sleep_seconds)

    maintainers = _normalize_owners(profile.get("owners"))
    summary = profile.get("summary") or record.repodata_summary or ""
    description = profile.get("description") or summary
    latest_version = profile.get("latest_version") or record.latest_version
    downloads = record.downloads if record.downloads >= 0 else extract_download_count(profile)

    return {
        "tier": tier,
        "domain": domain or "",
        "software_name": software_name,
        "package_name": record.package_name,
        "latest_version": latest_version,
        "dependencies": record.depends,
        "maintainers": maintainers,
        "description": description,
        "summary": summary,
        "downloads": downloads,
        "license": profile.get("license") or "",
        "home_url": profile.get("home") or "",
        "doc_url": profile.get("doc_url") or "",
        "dev_url": profile.get("dev_url") or "",
        "source_subdir": record.subdir,
        "timestamp": record.timestamp,
    }


def export_json(data: List[Dict], output_file: Path) -> None:
    output_file.parent.mkdir(parents=True, exist_ok=True)
    with output_file.open("w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)


def run_t0(
    index: Dict[str, PackageRecord],
    output_dir: Path,
    k_per_domain: int,
    k_overall: int,
    profile_cache: Dict[str, Dict],
) -> List[Dict]:
    result: List[Dict] = []
    for software_name, record, domain in select_t0_records(index, k_per_domain=k_per_domain, k_overall=k_overall):
        result.append(
            enrich_record(
                software_name=software_name,
                record=record,
                tier="T0",
                domain=domain,
                profile_cache=profile_cache,
            )
        )
    export_json(result, output_dir / "bioconda_t0_core_tools.json")
    return result


def run_t1(
    index: Dict[str, PackageRecord],
    output_dir: Path,
    max_per_domain: int,
    profile_cache: Dict[str, Dict],
    target_total: int,
) -> List[Dict]:
    result: List[Dict] = []
    selected = select_t1_records(index=index, max_per_domain=max_per_domain)
    selected = backfill_t1_records(index=index, selected=selected, target_total=target_total)
    for software_name, record, domain in selected:
        result.append(
            enrich_record(
                software_name=software_name,
                record=record,
                tier="T1",
                domain=domain,
                profile_cache=profile_cache,
            )
        )
    export_json(result, output_dir / "bioconda_t1_domain_tools.json")
    return result


def run_t2(
    index: Dict[str, PackageRecord],
    output_dir: Path,
    t2_tools: List[str],
    profile_cache: Dict[str, Dict],
) -> List[Dict]:
    result: List[Dict] = []
    selected = select_t2_records(index=index, tools=t2_tools)
    for software_name, record in selected:
        result.append(
            enrich_record(
                software_name=software_name,
                record=record,
                tier="T2",
                profile_cache=profile_cache,
            )
        )
    export_json(result, output_dir / "bioconda_t2_on_demand_tools.json")
    return result


def run_all_package_index(index: Dict[str, PackageRecord], output_dir: Path) -> None:
    rows = []
    for package_name, record in sorted(index.items()):
        rows.append(
            {
                "software_name": package_name,
                "package_name": package_name,
                "latest_version": record.latest_version,
                "dependencies": record.depends,
                "maintainers": [],
                "description": record.repodata_summary,
                "summary": record.repodata_summary,
                "downloads": record.downloads,
                "license": "",
                "home_url": "",
                "doc_url": "",
                "dev_url": "",
                "source_subdir": record.subdir,
                "timestamp": record.timestamp,
            }
        )
    export_json(rows, output_dir / "bioconda_all_packages_index.json")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Crawl Bioconda package metadata in a tiered manner (T0/T1/T2)."
    )
    parser.add_argument(
        "--tiers",
        nargs="+",
        default=["T0", "T1", "T2"],
        choices=["T0", "T1", "T2", "ALL"],
        help="Which tiers to run. Use ALL to export whole package index from repodata.",
    )
    parser.add_argument(
        "--subdirs",
        nargs="+",
        default=["linux-64", "noarch"],
        help="Bioconda subdirs to index, e.g. linux-64 noarch osx-64.",
    )
    parser.add_argument(
        "--output-dir",
        default="./output",
        help="Directory for output JSON files.",
    )
    parser.add_argument(
        "--t2-tools",
        nargs="*",
        default=[],
        help="Tools for T2 on-demand crawling, e.g. --t2-tools cellranger scanpy",
    )
    parser.add_argument(
        "--t1-max-per-domain",
        type=int,
        default=200,
        help="Max selected packages per T1 domain.",
    )
    parser.add_argument(
        "--t1-target-total",
        type=int,
        default=600,
        help="Target total rows for T1. Uses overall backfill when domain matches are insufficient.",
    )
    parser.add_argument(
        "--downloads-source",
        choices=["channel_api", "package_api"],
        default="channel_api",
        help="Download ranking source. channel_api first, package_api is fallback or explicit mode.",
    )
    parser.add_argument(
        "--t0-k-per-domain",
        type=int,
        default=5,
        help="T0: top K packages per domain ranked by downloads.",
    )
    parser.add_argument(
        "--t0-k-overall",
        type=int,
        default=20,
        help="T0: top K packages in overall ranking by downloads.",
    )
    parser.add_argument(
        "--package-api-max",
        type=int,
        default=0,
        help="Max packages for package_api mode (0 means all indexed packages).",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    print(f"[1/4] Building Bioconda latest-package index from subdirs: {args.subdirs}")
    latest_index = build_latest_index(subdirs=args.subdirs)
    print(f"      Indexed packages: {len(latest_index)}")
    profile_cache: Dict[str, Dict] = {}

    print("[2/4] Resolving package download counts...")
    download_index = resolve_download_index(
        index=latest_index,
        download_source=args.downloads_source,
        profile_cache=profile_cache,
        package_api_max=args.package_api_max,
    )
    for package_name, record in latest_index.items():
        record.downloads = download_index.get(package_name, -1)
    known_downloads = sum(1 for r in latest_index.values() if r.downloads >= 0)
    print(f"      Packages with download count: {known_downloads}")

    run_all = "ALL" in args.tiers
    if run_all:
        print("[3/4] Exporting full package index...")
        run_all_package_index(index=latest_index, output_dir=output_dir)
        print("      -> bioconda_all_packages_index.json")

    if "T0" in args.tiers:
        print("[4/4] Running T0 core tools (auto top-k by downloads)...")
        t0_rows = run_t0(
            index=latest_index,
            output_dir=output_dir,
            k_per_domain=args.t0_k_per_domain,
            k_overall=args.t0_k_overall,
            profile_cache=profile_cache,
        )
        print(f"      -> bioconda_t0_core_tools.json ({len(t0_rows)} rows)")

    if "T1" in args.tiers:
        print("[4/4] Running T1 domain tools...")
        t1_rows = run_t1(
            index=latest_index,
            output_dir=output_dir,
            max_per_domain=args.t1_max_per_domain,
            profile_cache=profile_cache,
            target_total=args.t1_target_total,
        )
        print(f"      -> bioconda_t1_domain_tools.json ({len(t1_rows)} rows)")

    if "T2" in args.tiers:
        print("[4/4] Running T2 on-demand tools...")
        t2_rows = run_t2(
            index=latest_index,
            output_dir=output_dir,
            t2_tools=args.t2_tools,
            profile_cache=profile_cache,
        )
        print(f"      -> bioconda_t2_on_demand_tools.json ({len(t2_rows)} rows)")

    print("Done.")


if __name__ == "__main__":
    main()