# ============================================================
# 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
# ============================================================

"""Download 10-K / 10-Q / 8-K primary documents for a list of companies.

    python scripts/download_filings.py --universe            # every company in universe.json
    python scripts/download_filings.py --ciks-file data/metadata/shards/M2U64.self.txt --rate 0.9
    python scripts/download_filings.py --ciks 0000320193,0000789019 --raw-dir /tmp/raw --log /tmp/log.jsonl

Resumable: already-logged filings whose file exists are skipped.
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

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

from src.acquisition.bulk_downloader import BulkDownloader
from src.acquisition.universe import load_universe
from src.utils.constants import DATA_DIRS, SEC_FORM_TYPES
from src.utils.logging import setup_logging

logger = setup_logging("download_filings")


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Bulk EDGAR downloader")
    g = p.add_mutually_exclusive_group(required=True)
    g.add_argument("--universe", action="store_true", help="all companies in data/metadata/universe.json")
    g.add_argument("--ciks-file", type=str, help="text file with one CIK per line")
    g.add_argument("--ciks", type=str, help="comma-separated CIKs")
    p.add_argument("--forms", type=str, default=",".join(SEC_FORM_TYPES))
    p.add_argument("--start", type=str, default="2010-01-01")
    p.add_argument("--rate", type=float, default=9.0, help="requests per second for this process")
    p.add_argument("--workers", type=int, default=8)
    p.add_argument("--raw-dir", type=str, default=None)
    p.add_argument("--log", type=str, default=None, help="download log (JSONL)")
    p.add_argument("--submissions-dir", type=str, default=None)
    p.add_argument("--known-logs", type=str, default=None,
                   help="comma-separated earlier download logs whose filings are skipped without re-download")
    return p.parse_args()


def main() -> None:
    args = parse_args()
    universe = {c["cik"]: c for c in load_universe()}
    if args.universe:
        ciks = list(universe)
    elif args.ciks_file:
        ciks = [l.strip() for l in Path(args.ciks_file).read_text().splitlines() if l.strip()]
    else:
        ciks = [c.strip() for c in args.ciks.split(",") if c.strip()]
    companies = [universe.get(c.zfill(10), {"cik": c.zfill(10), "ticker": ""}) for c in ciks]
    if not companies:
        logger.info("nothing to do")
        return
    dl = BulkDownloader(
        raw_dir=args.raw_dir, log_path=args.log, submissions_dir=args.submissions_dir,
        rate=args.rate, workers=args.workers,
        trusted_logs=[Path(x) for x in args.known_logs.split(",")] if args.known_logs else None,
    )
    stats = dl.download_many(companies, [f.strip() for f in args.forms.split(",")], args.start)
    logger.info("DONE %s", stats)
    if stats["failed"] > 0:
        sys.exit(2)


if __name__ == "__main__":
    main()