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

"""Merge the worker shards into the master database (run on the master node).

Expected layout (pulled by the orchestrator):
    data/shards/<node>/export/{companies,filings,filing_sections,filing_chunks,embeddings_metadata}.parquet
    data/shards/<node>/embeddings/model_name=<model>/*.parquet

Steps
 1. back up the master DuckDB, then INSERT OR REPLACE every shard table into it;
 2. move the shards' embedding files under data/parquet/embeddings/model_name=<model>/;
 3. rewrite legacy float64 embedding files as float32 (uniform schema);
 4. regenerate the single-file parquet exports used by SPID and build_shares
    (companies, filings, filing_sections, filing_chunks).
"""

from __future__ import annotations

import argparse
import shutil
import sys
from datetime import date
from pathlib import Path

import pyarrow as pa
import pyarrow.parquet as pq

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

from src.storage.database import DatabaseManager
from src.utils.constants import DATA_DIRS, DUCKDB_PATH, PARQUET_COMPRESSION, PARQUET_COMPRESSION_LEVEL
from src.utils.logging import setup_logging

logger = setup_logging("merge_shards")

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


PK = {"filings": "accession_number", "filing_sections": "section_id",
      "filing_chunks": "chunk_id", "embeddings_metadata": "embedding_id"}


def merge_tables(db: DatabaseManager, shard_dirs: list[Path]) -> None:
    """Append the shards' rows that the master does not have yet (keys are globally unique).

    ``companies`` is already complete on the master (same universe everywhere); only the
    ``shard_node`` tag is updated from the shard's filings.
    """
    for t in ["filings", "filing_sections", "filing_chunks", "embeddings_metadata"]:
        cols = [r[0] for r in db.execute(f"DESCRIBE {t}").fetchall()]
        pk = PK[t]
        for sd in shard_dirs:
            f = sd / "export" / f"{t}.parquet"
            if not f.exists():
                logger.warning("%s: missing %s", sd.name, f.name)
                continue
            file_cols = [c for c in pq.read_schema(str(f)).names if c in cols]
            col_list = ", ".join(file_cols)
            before = db.execute(f"SELECT COUNT(*) FROM {t}").fetchone()[0]
            db.execute(
                f"INSERT INTO {t} ({col_list}) SELECT {col_list} FROM read_parquet('{f}') s "
                f"WHERE NOT EXISTS (SELECT 1 FROM {t} m WHERE m.{pk} = s.{pk})"
            )
            after = db.execute(f"SELECT COUNT(*) FROM {t}").fetchone()[0]
            logger.info("%-20s <- %-8s +%d rows (now %d)", t, sd.name, after - before, after)
    for sd in shard_dirs:
        f = sd / "export" / "filings.parquet"
        if f.exists():
            db.execute(f"UPDATE companies SET shard_node = '{sd.name}' WHERE cik IN "
                       f"(SELECT DISTINCT cik FROM read_parquet('{f}'))")
    logger.info("companies.shard_node updated from the shards")


def move_embeddings(shard_dirs: list[Path]) -> None:
    dest_root = DATA_DIRS["parquet"] / "embeddings"
    for sd in shard_dirs:
        src_root = sd / "embeddings"
        if not src_root.exists():
            continue
        n = 0
        for part in src_root.glob("model_name=*"):
            dest = dest_root / part.name
            dest.mkdir(parents=True, exist_ok=True)
            for f in part.glob("*.parquet"):
                target = dest / f"{sd.name}-{f.name}"
                if not target.exists():
                    shutil.move(str(f), str(target))
                    n += 1
        logger.info("embeddings <- %-8s %d files", sd.name, n)


def normalise_legacy_embeddings() -> None:
    """Rewrite list<double> files as list<float> and drop the in-file model_name column."""
    root = DATA_DIRS["parquet"] / "embeddings"
    n = 0
    for f in root.glob("model_name=*/*.parquet"):
        schema = pq.read_schema(str(f))
        vec = schema.field("vector").type
        needs = pa.types.is_list(vec) and pa.types.is_float64(vec.value_type)
        needs = needs or ("model_name" in schema.names)
        if not needs:
            continue
        t = pq.read_table(str(f))
        if "model_name" in t.column_names:
            t = t.drop_columns(["model_name"])
        t = t.set_column(t.schema.get_field_index("vector"), "vector", t.column("vector").cast(pa.list_(pa.float32())))
        tmp = f.with_suffix(".tmp.parquet")
        pq.write_table(t, str(tmp), compression=PARQUET_COMPRESSION, compression_level=PARQUET_COMPRESSION_LEVEL)
        tmp.replace(f)
        n += 1
    logger.info("normalised %d legacy embedding files to float32", n)


def export_single_parquet(db: DatabaseManager) -> None:
    for t in ["companies", "filings", "filing_sections", "filing_chunks"]:
        out_dir = DATA_DIRS["parquet"] / t
        out_dir.mkdir(parents=True, exist_ok=True)
        # remove stale hive partitions from earlier runs so the directory is one clean file
        for stale in out_dir.glob("form_type=*"):
            shutil.rmtree(stale, ignore_errors=True)
        path = out_dir / f"{t}.parquet"
        db.execute(f"COPY (SELECT * FROM {t}) TO '{path}' (FORMAT PARQUET, COMPRESSION ZSTD, ROW_GROUP_SIZE 200000)")
        logger.info("parquet %-18s -> %s (%.1f MB)", t, path, path.stat().st_size / 1e6)


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--shards-dir", type=str, default=str(DATA_DIRS["raw"].parent.parent / "shards"))
    p.add_argument("--nodes", type=str, default=None, help="comma list; default = every sub-directory")
    p.add_argument("--no-backup", action="store_true")
    args = p.parse_args()

    shards_root = Path(args.shards_dir)
    shard_dirs = sorted(d for d in shards_root.iterdir() if d.is_dir()) if shards_root.exists() else []
    if args.nodes:
        wanted = {n.strip() for n in args.nodes.split(",")}
        shard_dirs = [d for d in shard_dirs if d.name in wanted]
    logger.info("merging %d shard(s): %s", len(shard_dirs), [d.name for d in shard_dirs])

    if not args.no_backup and DUCKDB_PATH.exists():
        bak = DUCKDB_PATH.with_name(f"{DUCKDB_PATH.stem}.pre-merge-{date.today():%Y%m%d}.duckdb")
        if not bak.exists():
            shutil.copy2(DUCKDB_PATH, bak)
            logger.info("backup -> %s", bak)

    db = DatabaseManager()
    db.initialize()
    db.execute("SET preserve_insertion_order=false")
    merge_tables(db, shard_dirs)
    move_embeddings(shard_dirs)
    normalise_legacy_embeddings()
    export_single_parquet(db)
    for t, n in db.get_stats().items():
        logger.info("  %-25s %d rows", t, n)
    db.close()
    logger.info("merge complete")


if __name__ == "__main__":
    main()