"""Seed / update the `disciplines` reference collection.

Each document maps an arXiv discipline to its human name and a PHIDESK topic:

    { source: "arxiv", fullname: "Computer Science", shortname: "cs",
      phidesktopic: "" }

`phidesktopic` is left blank on first insert for you to fill in; re-running this
script never overwrites a value you have already set (it only refreshes the
fullname and keeps everything else via $setOnInsert).

Run from the API folder:  python seed_disciplines.py
"""

from pymongo import MongoClient, UpdateOne

from app.config import settings

COLLECTION = "disciplines"

# The arXiv top-level groups currently present in the collected data. Add rows
# here (or straight into the collection) to cover more.
ARXIV_DISCIPLINES = [
    ("cs", "Computer Science"),
    ("econ", "Economics"),
    ("eess", "Electrical Engineering and Systems Science"),
    ("math", "Mathematics"),
    ("grp_physics", "Physics"),
    ("q-bio", "Quantitative Biology"),
    ("q-fin", "Quantitative Finance"),
    ("stat", "Statistics"),
]


def main() -> None:
    client = MongoClient(settings.mongo_uri)
    col = client[settings.mongo_db][COLLECTION]

    # A discipline is unique per (source, shortname).
    col.create_index([("source", 1), ("shortname", 1)], unique=True)

    ops = [
        UpdateOne(
            {"source": "arxiv", "shortname": short},
            {
                # Always keep the name current...
                "$set": {"source": "arxiv", "shortname": short, "fullname": full},
                # ...but never clobber a phidesktopic you have filled in.
                "$setOnInsert": {"phidesktopic": ""},
            },
            upsert=True,
        )
        for short, full in ARXIV_DISCIPLINES
    ]
    result = col.bulk_write(ops)
    print(f"disciplines: upserted {result.upserted_count}, matched {result.matched_count}")
    print(f"total in '{COLLECTION}': {col.count_documents({})}")
    client.close()


if __name__ == "__main__":
    main()