"""
Migration plan step 1: raw staging load.

Dumps every non-scratch table from all four source Access files 1:1 into
either MySQL (one database per source system, per config.SOURCES — e.g.
`stg_utilities`, `stg_seguros`, `stg_scothia`, matching a MySQL "schema" to
a MySQL "database" 1:1) or, if --output-dir is given, to Parquet files —
useful for environments (like this one) that have Access + pyodbc available
but not a live MySQL instance reachable to load into yet.

Every staged table gets two extra columns: _legacy_source_table (the
original Access table name) and _row_num (ordinal position in the source
table) so rows are traceable even before any real key is identified.

Usage:
    python load_staging.py --output-dir ./output                       # Parquet, no DB needed
    python load_staging.py --database-url mysql+pymysql://user:pass@host/db  # load into MySQL
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

import pandas as pd

from config import SOURCES
import extract


def stage_source(source_name: str, source_cfg: dict, sink) -> None:
    path = source_cfg["path"]
    if not Path(path).exists():
        print(f"  [skip] {source_name}: file not found at {path}", file=sys.stderr)
        return

    print(f"=== {source_name} ({path}) ===")
    cnxn = extract.connect(path)
    tables = extract.list_tables(cnxn)
    excluded = source_cfg["exclude"]

    for table_name in tables:
        if table_name in excluded:
            print(f"  [exclude] {table_name}")
            continue

        try:
            df = extract.read_table(cnxn, table_name)
        except Exception as exc:  # noqa: BLE001 - report and keep going
            print(f"  [ERROR] {table_name}: {exc}", file=sys.stderr)
            continue

        df.insert(0, "_legacy_source_table", table_name)
        df.insert(1, "_row_num", range(len(df)))

        sink(source_cfg["schema"], table_name, df)
        print(f"  [ok] {table_name}: {len(df)} rows, {len(df.columns) - 2} columns")


def make_parquet_sink(output_dir: Path):
    def sink(schema: str, table_name: str, df: pd.DataFrame) -> None:
        target_dir = output_dir / schema
        target_dir.mkdir(parents=True, exist_ok=True)
        safe_name = extract.sanitize_column_name(table_name)
        df.to_parquet(target_dir / f"{safe_name}.parquet", index=False)

    return sink


def make_mysql_sink(database_url: str):
    from sqlalchemy import create_engine, text

    # A MySQL "schema" is a database - CREATE SCHEMA is a synonym for
    # CREATE DATABASE. Connect without a default database in the URL so
    # each staging area (stg_utilities/stg_seguros/stg_scothia) can be
    # created and written to independently by one engine.
    engine = create_engine(database_url)
    created_schemas: set[str] = set()

    def sink(schema: str, table_name: str, df: pd.DataFrame) -> None:
        if schema not in created_schemas:
            with engine.begin() as conn:
                conn.execute(text(f"CREATE SCHEMA IF NOT EXISTS `{schema}`"))
            created_schemas.add(schema)

        safe_name = extract.sanitize_column_name(table_name)
        df.to_sql(
            safe_name,
            engine,
            schema=schema,
            if_exists="replace",
            index=False,
        )

    return sink


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output-dir", type=Path, help="Write Parquet files here instead of a database")
    parser.add_argument("--database-url", type=str, help="MySQL connection string, e.g. mysql+pymysql://user:pass@host:3306/ (sqlalchemy format)")
    parser.add_argument("--only", type=str, help="Comma-separated subset of source names to run (utilities,seguros,scothia)")
    args = parser.parse_args()

    if not args.output_dir and not args.database_url:
        parser.error("one of --output-dir or --database-url is required")

    sink = make_parquet_sink(args.output_dir) if args.output_dir else make_mysql_sink(args.database_url)

    only = set(args.only.split(",")) if args.only else None

    for source_name, source_cfg in SOURCES.items():
        if only and source_name not in only:
            continue
        stage_source(source_name, source_cfg, sink)


if __name__ == "__main__":
    main()
