#!/usr/bin/env python
from __future__ import annotations

import argparse
from pathlib import Path

from assaylens.data.access import iter_limited_rows
from assaylens.data.hf_upload import auto_upload_folder
from assaylens.data.schemas import DATASET_QUEST_HF_ID, GOA_HF_ID
from assaylens.features.goa import build_goa_count_features
from assaylens.reports.markdown import write_markdown


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--limit", type=int, default=None)
    parser.add_argument("--out-dir", default="reports")
    args = parser.parse_args()

    goa_rows = list(iter_limited_rows(GOA_HF_ID, "train", args.limit))
    quest_rows = list(iter_limited_rows(DATASET_QUEST_HF_ID, "train", args.limit))
    goa_features = build_goa_count_features(goa_rows)

    write_markdown(
        Path(args.out_dir) / "source_inventory.md",
        "AssayLens Source Inventory",
        [
            (
                "Used Sources",
                "\n".join(
                    [
                        "- core dataset: `eve-bio/drug-target-activity`",
                        "- Hugging Science support datasets: `goa_reasoning`, `dataset-quest-index`",
                        f"- GOA rows inspected: {len(goa_rows)}",
                        f"- Dataset Quest rows inspected: {len(quest_rows)}",
                        f"- GOA annotation keys in inspected rows: {len(goa_features)}",
                    ]
                ),
            ),
            (
                "Excluded Hugging Science Datasets",
                "\n".join(
                    [
                        "- antibody datasets: out of scope for small-molecule assay measurement model",
                        "- AI4FA microbiome: out of scope",
                        "- ARC/ISLES/AOMIC imaging: out of scope",
                        "- m-boltz-submissions: too small and not relevant to v1",
                    ]
                ),
            ),
        ],
    )
    uploaded = auto_upload_folder(args.out_dir, "reports")
    if uploaded:
        print(f"uploaded Hugging Science report to {uploaded}/reports")
    print(f"wrote {Path(args.out_dir) / 'source_inventory.md'}")


if __name__ == "__main__":
    main()