#!/usr/bin/env python3
"""Rerun this card: who is hiring to build and run US data centers.

    export METIX_KEY=metix_xxxxxxxxxxxx   (or put it in the repository's .env)
    python3 cases/data-center-jobs-us-2026/fetch.py

Counts only, one API Credit each, about 75 in all; no posting is read.

  states.json     data center job titles by state, the rest as one row
  employers.json  the same postings by employer
  roles.json      the same postings by the word the title uses for the job
  trades.json     electricians and construction leads whose posting names data centers
"""

from __future__ import annotations

import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE.parents[1] / "tools"))

from metix_client import Platform, aggregate, load_query, write_json

# Every state with 50 or more data center titles in a first pass, largest first after
# the run; the rest (and postings with no state) are one row.
STATES = [
    "Texas", "Virginia", "California", "Ohio", "Georgia", "North Carolina", "Oregon",
    "Washington", "Arizona", "New York", "Indiana", "Pennsylvania", "Illinois",
    "Colorado", "Nevada", "Florida", "Tennessee", "Iowa", "Missouri", "New Jersey",
    "Maryland", "Mississippi", "South Carolina", "Louisiana", "Oklahoma", "Utah",
    "Wisconsin", "Minnesota", "Massachusetts", "Nebraska",
]
# Employer names exactly as the records spell them; Amazon posts under two.
EMPLOYERS = {
    "Amazon": ["Amazon Web Services (AWS)", "Amazon"],
    "Oracle": ["Oracle"],
    "Google": ["Google"],
    "Meta": ["Meta"],
    "QTS Data Centers": ["QTS Data Centers"],
    "CoreWeave": ["CoreWeave"],
    "Serverfarm": ["Serverfarm"],
    "Microsoft": ["Microsoft"],
    "Equinix": ["Equinix"],
    "Switch": ["Switch"],
}
# The word a title uses for the job. Titles can hold more than one, so rows overlap.
ROLES = ["technician", "manager", "engineer", "construction", "operations", "mechanical", "security"]
# States where the trades comparison is made: the largest data center states, plus Arizona.
TRADE_STATES = ["Texas", "Virginia", "Arizona", "Georgia", "Ohio", "California"]


def main() -> int:
    platform = Platform(max_credits=90)
    titles = load_query(HERE, "data-center-titles")["where"]
    mentions = load_query(HERE, "mentions-data-centers")["where"]
    electricians = load_query(HERE, "electricians")["where"]
    leads = load_query(HERE, "construction-leads")["where"]

    def count(*conds: dict) -> int:
        n = platform.count("jobs", {"all": list(conds)})
        assert isinstance(n, int), n
        return n

    def state(s: str) -> dict:
        return {"field": "location.state", "eq": s}

    # States: every listed state, then everything else in one count, so the rows add
    # up to the national total without asking for a total that may be banded.
    rows = [{"group": s, "count": count(titles, state(s))} for s in STATES]
    rest = count(titles, {"not": [{"field": "location.state", "in": STATES}]})
    total = sum(r["count"] for r in rows) + rest
    rows.sort(key=lambda r: -r["count"])
    for r in rows:
        r["share"] = round(r["count"] / total, 4)
    rows.append({"group": "other", "count": rest, "share": round(rest / total, 4)})
    write_json(HERE, "states.json", aggregate("jobs", platform.snapshot, "queries/data-center-titles.json", rows, us_total=total))

    amazon = {"field": "company.name", "in": EMPLOYERS["Amazon"]}
    emp = []
    for name, spellings in EMPLOYERS.items():
        n = count(titles, {"field": "company.name", "in": spellings})
        emp.append({"group": name, "count": n, "share": round(n / total, 4)})
    emp.sort(key=lambda r: -r["count"])
    by_state = {s: count(titles, amazon, state(s)) for s in ["Texas", "Virginia"]}
    write_json(HERE, "employers.json", aggregate("jobs", platform.snapshot, "queries/data-center-titles.json", emp, us_total=total, amazon_by_state=by_state))

    roles = [{"group": w, "count": count(titles, {"field": "title", "match": w})} for w in ROLES]
    for r in roles:
        r["share"] = round(r["count"] / total, 4)
    write_json(HERE, "roles.json", aggregate("jobs", platform.snapshot, "queries/data-center-titles.json", roles, us_total=total, overlapping=True))

    trades = []
    for label, pop in [("electricians", electricians), ("construction leads", leads)]:
        all_n, dc_n = count(pop), count(pop, mentions)
        trades.append({"group": label, "state": "United States", "count": dc_n, "of": all_n, "share": round(dc_n / all_n, 4)})
    for s in TRADE_STATES:
        all_n, dc_n = count(electricians, state(s)), count(electricians, state(s), mentions)
        trades.append({"group": "electricians", "state": s, "count": dc_n, "of": all_n, "share": round(dc_n / all_n, 4)})
    for s in ["Texas", "Virginia"]:
        dc_n = count(leads, state(s), mentions)
        trades.append({"group": "construction leads", "state": s, "count": dc_n})
    write_json(HERE, "trades.json", aggregate("jobs", platform.snapshot, "queries/electricians.json", trades))

    write_json(HERE, "receipt.json", platform.receipt())
    print(total, rows[:3], f"{platform.spent()} API Credits")
    return 0


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