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

"""Concurrent, rate-limited, resumable EDGAR downloader (2026-10-27).

Same on-disk layout and download-log format as ``FilingDownloader`` so that
``scripts/build_database.py`` works unchanged, but:

* one global token bucket shared by all threads (SEC allows 10 req/s per IP;
  several cluster nodes behind the same IP each run with ``rate = 9 / n``);
* a thread pool keeps the pipe full (latency no longer caps throughput);
* HTTP 403/429 (rate-limit ban) triggers a cool-down instead of a hard failure;
* the trimmed company header of every submissions JSON (SIC, state, exchanges)
  is saved to ``metadata/submissions/<cik>.json`` for the companies table.
"""

from __future__ import annotations

import json
import os
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone
from pathlib import Path

import httpx

from src.utils.constants import DATA_DIRS, SEC_ARCHIVES_URL, SEC_SUBMISSIONS_URL
from src.utils.logging import get_logger

logger = get_logger("acquisition.bulk")

MAX_RETRIES = 6
COOLDOWN_SECONDS = 65.0


class TokenBucket:
    def __init__(self, rate: float, capacity: float | None = None) -> None:
        self.rate = max(rate, 0.05)
        self.capacity = capacity or max(1.0, self.rate)
        self._tokens = self.capacity
        self._last = time.monotonic()
        self._lock = threading.Lock()

    def acquire(self) -> None:
        while True:
            with self._lock:
                now = time.monotonic()
                self._tokens = min(self.capacity, self._tokens + (now - self._last) * self.rate)
                self._last = now
                if self._tokens >= 1.0:
                    self._tokens -= 1.0
                    return
                wait = (1.0 - self._tokens) / self.rate
            time.sleep(wait)


class BulkDownloader:
    def __init__(
        self,
        raw_dir: Path | None = None,
        log_path: Path | None = None,
        submissions_dir: Path | None = None,
        rate: float = 9.0,
        workers: int = 8,
        user_agent: str | None = None,
        trusted_logs: list[Path] | None = None,
    ) -> None:
        self.raw_dir = Path(raw_dir) if raw_dir else DATA_DIRS["raw"]
        self.log_path = Path(log_path) if log_path else DATA_DIRS["metadata"] / "download_log.json"
        self.submissions_dir = Path(submissions_dir) if submissions_dir else DATA_DIRS["metadata"] / "submissions"
        self.bucket = TokenBucket(rate)
        self.workers = workers
        self.user_agent = user_agent or os.environ.get("SEC_USER_AGENT", "UQO-Research research@uqo.ca")
        self._client = httpx.Client(
            timeout=60,
            headers={"User-Agent": self.user_agent, "Accept-Encoding": "gzip, deflate"},
            follow_redirects=True,
            limits=httpx.Limits(max_connections=workers + 4, max_keepalive_connections=workers + 4),
        )
        self._log_lock = threading.Lock()
        self._cooldown_until = 0.0
        self._cooldown_lock = threading.Lock()
        self._done_keys: set[str] = set()
        self._trusted_keys: set[str] = set()
        self._load_log()
        for tl in trusted_logs or []:
            self._load_trusted(Path(tl))
        self.stats = {"ciks": 0, "filings_listed": 0, "downloaded": 0, "skipped": 0, "failed": 0, "bytes": 0}
        self._stats_lock = threading.Lock()

    # ---------------------------------------------------------------- HTTP
    def _get(self, url: str, expect_json: bool = False):
        for attempt in range(MAX_RETRIES):
            with self._cooldown_lock:
                until = self._cooldown_until
            now = time.monotonic()
            if until > now:
                time.sleep(until - now)
            self.bucket.acquire()
            try:
                r = self._client.get(url)
            except httpx.RequestError as exc:
                wait = min(60.0, 2.0 ** attempt)
                logger.warning("request error %s (%s) attempt %d, sleeping %.0fs", url, exc, attempt + 1, wait)
                time.sleep(wait)
                continue
            if r.status_code == 404:
                return None
            if r.status_code in (403, 429, 503):
                with self._cooldown_lock:
                    self._cooldown_until = time.monotonic() + COOLDOWN_SECONDS
                logger.warning("HTTP %d on %s -> cooling down %.0fs (attempt %d)",
                               r.status_code, url, COOLDOWN_SECONDS, attempt + 1)
                continue
            if r.status_code >= 400:
                wait = min(60.0, 2.0 ** attempt)
                logger.warning("HTTP %d on %s attempt %d, sleeping %.0fs", r.status_code, url, attempt + 1, wait)
                time.sleep(wait)
                continue
            return r.json() if expect_json else r.content
        logger.error("giving up on %s", url)
        return None

    # ------------------------------------------------------------- listing
    def list_filings(self, cik: str, forms: list[str], start_date: str) -> tuple[dict | None, list[dict]]:
        padded = cik.zfill(10)
        data = self._get(f"{SEC_SUBMISSIONS_URL}/CIK{padded}.json", expect_json=True)
        if data is None:
            return None, []
        header = {
            "cik": padded,
            "name": data.get("name"),
            "sic": data.get("sic"),
            "sic_description": data.get("sicDescription"),
            "state_of_incorporation": data.get("stateOfIncorporation"),
            "state_of_business": (data.get("addresses", {}).get("business", {}) or {}).get("stateOrCountry"),
            "fiscal_year_end": data.get("fiscalYearEnd"),
            "exchanges": data.get("exchanges"),
            "tickers": data.get("tickers"),
            "entity_type": data.get("entityType"),
            "former_names": [n.get("name") for n in data.get("formerNames", [])],
            "fetched_at": datetime.now(timezone.utc).isoformat(),
        }
        blocks = [data.get("filings", {}).get("recent", {})]
        for f in data.get("filings", {}).get("files", []):
            name = f.get("name")
            if not name:
                continue
            # older blocks are needed only if they can contain filings after start_date
            if f.get("filingTo", "9999") < start_date:
                continue
            older = self._get(f"{SEC_SUBMISSIONS_URL}/{name}", expect_json=True)
            if older:
                blocks.append(older)
        filings: list[dict] = []
        wanted = set(forms)
        for b in blocks:
            fm, acc, fd, rd, pd_ = (b.get("form", []), b.get("accessionNumber", []), b.get("filingDate", []),
                                    b.get("reportDate", []), b.get("primaryDocument", []))
            for i, form in enumerate(fm):
                if form not in wanted:
                    continue
                if i >= len(fd) or fd[i] < start_date:
                    continue
                primary = pd_[i] if i < len(pd_) else ""
                if not primary:
                    continue
                accession = acc[i]
                filings.append({
                    "accession_number": accession,
                    "form_type": form,
                    "filing_date": fd[i],
                    "report_date": rd[i] if i < len(rd) else "",
                    "primary_document": primary,
                    "filing_url": f"{SEC_ARCHIVES_URL}/{padded.lstrip('0')}/{accession.replace('-', '')}/{primary}",
                })
        return header, filings

    # ------------------------------------------------------------ download
    def _local_path(self, cik: str, filing: dict) -> Path:
        return (self.raw_dir / cik.zfill(10) / filing["form_type"]
                / filing["accession_number"].replace("-", "") / filing["primary_document"])

    def download_cik(self, cik: str, ticker: str, forms: list[str], start_date: str) -> dict:
        cik = cik.zfill(10)
        header, filings = self.list_filings(cik, forms, start_date)
        res = {"cik": cik, "listed": len(filings), "downloaded": 0, "skipped": 0, "failed": 0}
        if header is None:
            logger.warning("no submissions for CIK %s (%s)", cik, ticker)
            res["failed"] = 1
            return res
        self.submissions_dir.mkdir(parents=True, exist_ok=True)
        (self.submissions_dir / f"{cik}.json").write_text(json.dumps(header, indent=1))
        if not ticker and header.get("tickers"):
            ticker = header["tickers"][0]

        for f in filings:
            key = f"{cik}|{f['accession_number']}|{f['primary_document']}"
            path = self._local_path(cik, f)
            if key in self._trusted_keys:
                res["skipped"] += 1
                continue
            if key in self._done_keys and path.exists():
                res["skipped"] += 1
                continue
            if path.exists() and path.stat().st_size > 0:
                self._record(key, cik, ticker, f, path, path.stat().st_size)
                res["skipped"] += 1
                continue
            content = self._get(f["filing_url"])
            if content is None:
                res["failed"] += 1
                continue
            path.parent.mkdir(parents=True, exist_ok=True)
            tmp = path.with_suffix(path.suffix + ".part")
            tmp.write_bytes(content)
            os.replace(tmp, path)
            self._record(key, cik, ticker, f, path, len(content))
            res["downloaded"] += 1
            with self._stats_lock:
                self.stats["bytes"] += len(content)
        return res

    def download_many(self, companies: list[dict], forms: list[str], start_date: str) -> dict:
        t0 = time.time()
        total = len(companies)
        logger.info("bulk download: %d companies, forms=%s, since %s, rate=%.2f/s, workers=%d, raw=%s",
                    total, forms, start_date, self.bucket.rate, self.workers, self.raw_dir)
        done = 0
        with ThreadPoolExecutor(max_workers=self.workers) as ex:
            futs = {ex.submit(self.download_cik, c["cik"], c.get("ticker", ""), forms, start_date): c
                    for c in companies}
            for fut in as_completed(futs):
                c = futs[fut]
                done += 1
                try:
                    r = fut.result()
                except Exception:  # noqa: BLE001
                    logger.exception("CIK %s failed", c["cik"])
                    with self._stats_lock:
                        self.stats["failed"] += 1
                    continue
                with self._stats_lock:
                    self.stats["ciks"] += 1
                    self.stats["filings_listed"] += r["listed"]
                    self.stats["downloaded"] += r["downloaded"]
                    self.stats["skipped"] += r["skipped"]
                    self.stats["failed"] += r["failed"]
                    s = dict(self.stats)
                el = time.time() - t0
                logger.info("[%d/%d] %s %s listed=%d dl=%d skip=%d fail=%d | total dl=%d skip=%d fail=%d %.1f GB %.0fs",
                            done, total, c["cik"], c.get("ticker", ""), r["listed"], r["downloaded"],
                            r["skipped"], r["failed"], s["downloaded"], s["skipped"], s["failed"],
                            s["bytes"] / 1e9, el)
        self._client.close()
        logger.info("bulk download complete: %s in %.0fs", self.stats, time.time() - t0)
        return dict(self.stats)

    # ------------------------------------------------------------------ log
    def _load_log(self) -> None:
        if not self.log_path.exists():
            return
        with open(self.log_path) as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    key = json.loads(line).get("key", "")
                except json.JSONDecodeError:
                    continue
                if key:
                    self._done_keys.add(key)
        logger.info("download log: %d entries already done (%s)", len(self._done_keys), self.log_path)

    def _load_trusted(self, path: Path) -> None:
        """Keys from an earlier corpus (e.g. the S&P 500 log) are skipped even if the file is not here yet."""
        if not path.exists():
            logger.warning("trusted log %s not found", path)
            return
        n = 0
        with open(path) as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    key = json.loads(line).get("key", "")
                except json.JSONDecodeError:
                    continue
                if key:
                    self._trusted_keys.add(key)
                    n += 1
        logger.info("trusted log: %d keys from %s", n, path)

    def _record(self, key: str, cik: str, ticker: str, filing: dict, path: Path, size: int) -> None:
        entry = {
            "key": key,
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "ticker": ticker,
            "form_type": filing["form_type"],
            "accession_number": filing["accession_number"],
            "filing_date": filing["filing_date"],
            "report_date": filing.get("report_date", ""),
            "primary_document": filing["primary_document"],
            "filing_url": filing["filing_url"],
            "file_size_bytes": size,
            "local_path": str(path),
        }
        with self._log_lock:
            self._done_keys.add(key)
            self.log_path.parent.mkdir(parents=True, exist_ok=True)
            with open(self.log_path, "a") as f:
                f.write(json.dumps(entry) + "\n")