# ============================================================
# Project: EDGAR Research Database
# Description: SEC EDGAR Research Infrastructure
#
# Authors:
# - Manel Kammoun — manel.kammoun@uqo.ca
# - Djerrycharli Tandjambianda — djerrycharli.tandjambianda@uqo.ca
# - Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
#
# Institution:
# Université du Québec en Outaouais (UQO)
#
# License:
# Research / Academic Use
# ============================================================

"""Build data/metadata/universe.json = S&P 500 (cached) ∪ Russell 3000 (iShares IWV).

Optionally split the companies into per-node shards for the cluster run:

    python scripts/build_universe.py --shards M3U96a:28,M2U64:24,... --bhs64-share 0.5

writes data/metadata/shards.json plus data/metadata/shards/<node>.{self,bhs64}.txt
(CIK lists).  Companies already present locally (the S&P 500 corpus) are pinned
to --master and are not re-sharded.
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from src.acquisition.cik_mapper import CIKMapper
from src.acquisition.sp500 import SP500Universe
from src.acquisition.universe import Russell3000Universe, build_universe, load_universe
from src.utils.constants import DATA_DIRS
from src.utils.logging import setup_logging

logger = setup_logging("build_universe")


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Build the company universe (S&P 500 ∪ Russell 3000)")
    p.add_argument("--refresh-sp500", action="store_true", help="Re-scrape the S&P 500 list (default: cached)")
    p.add_argument("--skip-fetch", action="store_true", help="Reuse cached Russell 3000 file")
    p.add_argument("--shards", type=str, default=None,
                   help="node:cores,... for the new companies (excludes --master)")
    p.add_argument("--master", type=str, default="M3U96b", help="node holding the existing S&P 500 corpus")
    p.add_argument("--bhs64-share", type=float, default=0.5,
                   help="fraction of each node's new CIKs downloaded through BHS64 (second IP)")
    return p.parse_args()


def main() -> None:
    args = parse_args()
    sp = SP500Universe()
    sp500 = sp.get_constituents() if args.refresh_sp500 else sp.load_cached()

    mapper = CIKMapper()
    mapper.build_mapping()
    ru = Russell3000Universe()
    russell = ru.load_cached() if args.skip_fetch else ru.get_constituents(mapper)

    universe = build_universe(sp500, russell)

    if not args.shards:
        return

    existing = {str(c["cik"]).zfill(10) for c in sp500 if c.get("cik")}
    new_ciks = [e["cik"] for e in universe if e["cik"] not in existing]
    nodes = []
    for item in args.shards.split(","):
        name, cores = item.split(":")
        nodes.append((name.strip(), int(cores)))
    total = sum(c for _, c in nodes)

    # weight each node by cores; deal CIKs round-robin weighted so sizes are proportional
    shares = {n: c / total for n, c in nodes}
    assigned: dict[str, list[str]] = {n: [] for n, _ in nodes}
    assigned[args.master] = []
    quota = {n: shares[n] * len(new_ciks) for n, _ in nodes}
    # largest-remainder allocation on a stable CIK order
    order = sorted(new_ciks)
    idx = 0
    for n, _ in nodes:
        k = int(round(quota[n]))
        assigned[n].extend(order[idx:idx + k])
        idx += k
    for leftover in order[idx:]:
        assigned[nodes[0][0]].append(leftover)

    shard_dir = DATA_DIRS["metadata"] / "shards"
    shard_dir.mkdir(parents=True, exist_ok=True)
    manifest = {"master": args.master, "bhs64_share": args.bhs64_share, "nodes": {}}
    for n, ciks in assigned.items():
        ciks = sorted(ciks)
        n_bhs = int(round(len(ciks) * args.bhs64_share))
        bhs = ciks[:n_bhs]
        own = ciks[n_bhs:]
        if n == args.master:
            own = sorted(existing)  # refresh of the S&P 500 corpus on the master
            bhs = []
        (shard_dir / f"{n}.self.txt").write_text("\n".join(own) + ("\n" if own else ""))
        (shard_dir / f"{n}.bhs64.txt").write_text("\n".join(bhs) + ("\n" if bhs else ""))
        manifest["nodes"][n] = {"self": own, "bhs64": bhs, "n_total": len(own) + len(bhs)}
        logger.info("shard %-8s self=%4d bhs64=%4d", n, len(own), len(bhs))
    with open(DATA_DIRS["metadata"] / "shards.json", "w") as f:
        json.dump(manifest, f, indent=2)
    logger.info("Wrote %s (%d new CIKs over %d nodes)", DATA_DIRS["metadata"] / "shards.json",
                len(new_ciks), len(nodes))


if __name__ == "__main__":
    main()