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

"""Export this node's DuckDB tables to data/export/<table>.parquet for the merge.

Embeddings are already parquet files under data/parquet/embeddings/ and are
rsync'ed as-is by the orchestrator.
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

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

from src.storage.database import DatabaseManager
from src.utils.constants import PROJECT_ROOT
from src.utils.logging import setup_logging

logger = setup_logging("export_shard")

TABLES = ["companies", "filings", "filing_sections", "filing_chunks", "embeddings_metadata"]


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--out", type=str, default=str(PROJECT_ROOT / "data" / "export"))
    p.add_argument("--tables", type=str, default=",".join(TABLES))
    args = p.parse_args()
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    db = DatabaseManager()
    for t in args.tables.split(","):
        t = t.strip()
        path = out / f"{t}.parquet"
        n = db.execute(f"SELECT COUNT(*) FROM {t}").fetchone()[0]
        db.execute(f"COPY (SELECT * FROM {t}) TO '{path}' (FORMAT PARQUET, COMPRESSION ZSTD, ROW_GROUP_SIZE 200000)")
        logger.info("exported %-20s %10d rows -> %s (%.1f MB)", t, n, path, path.stat().st_size / 1e6)
    db.close()
    (out / "EXPORT_DONE").write_text("ok\n")


if __name__ == "__main__":
    main()