"""Load a BIRD SQLite dataset into Postgres.

Runs inside the api container (`make seed`). Reads which dataset to load
from `settings.dataset` (env var `NVI_DATASET`, default `financial`) and
expects the source SQLite at `ctrl/seed/data/<dataset>.sqlite` — produced
on the host by `ctrl/seed/download.sh` and reaching the container via
Tilt live_update sync (kind) or the docker-compose bind mount.

DDL is built from SQLAlchemy `MetaData` + `Table` + `Column` objects, not
string SQL. Bulk insert uses psycopg's COPY protocol via the engine's raw
connection — there's no SQLAlchemy equivalent, and INSERT for ~1M rows
would be orders of magnitude slower. Identifiers in the COPY statement
are rendered through SA's dialect-aware identifier preparer.
"""

from __future__ import annotations

import sqlite3
import sys
from pathlib import Path
from typing import Any

from sqlalchemy import (
    BigInteger, Boolean, Column, Date, DateTime, Float, Integer, LargeBinary,
    MetaData, Numeric, SmallInteger, Table, Text, text,
)
from sqlalchemy.engine import Engine

from api.config import get_settings
from api.tools.db import get_engine

# SQLite is type-affinity, not strict — pick a SQLAlchemy type per affinity.
# Anything unrecognised falls through to Text, which is always safe.
TYPE_MAP: dict[str, Any] = {
    "INT": BigInteger,
    "INTEGER": BigInteger,
    "TINYINT": SmallInteger,
    "SMALLINT": SmallInteger,
    "MEDIUMINT": Integer,
    "BIGINT": BigInteger,
    "REAL": Float,
    "DOUBLE": Float,
    "FLOAT": Float,
    "NUMERIC": Numeric,
    "DECIMAL": Numeric,
    "BOOLEAN": Boolean,
    "DATE": Date,
    "DATETIME": DateTime,
    "TIMESTAMP": DateTime,
    "TEXT": Text,
    "CHAR": Text,
    "VARCHAR": Text,
    "CLOB": Text,
    "BLOB": LargeBinary,
}


def _sa_type(sqlite_type: str):
    base = sqlite_type.upper().split("(")[0].strip()
    return TYPE_MAP.get(base, Text)


def _list_tables(src: sqlite3.Connection) -> list[str]:
    return [
        r[0] for r in src.execute(
            "SELECT name FROM sqlite_master "
            "WHERE type='table' AND name NOT LIKE 'sqlite_%' "
            "ORDER BY name"
        ).fetchall()
    ]


def _discover_table(src: sqlite3.Connection, metadata: MetaData,
                    name: str, schema: str) -> Table:
    cols = src.execute(f'PRAGMA table_info("{name}")').fetchall()
    # PRAGMA returns (cid, name, type, notnull, dflt_value, pk).
    sa_cols = [
        Column(
            c[1],
            _sa_type(c[2]),
            nullable=not c[3],
            primary_key=bool(c[5]),
        )
        for c in cols
    ]
    return Table(name, metadata, *sa_cols, schema=schema)


def _copy_table(engine: Engine, src: sqlite3.Connection, sa_table: Table) -> int:
    """Bulk-load via psycopg COPY. Identifiers go through SA's dialect preparer
    so we never hand-quote names."""
    prep = engine.dialect.identifier_preparer
    table_ident = prep.format_table(sa_table)                       # "<schema>"."<table>"
    col_list = ", ".join(prep.quote(c.name) for c in sa_table.columns)
    src_ident = '"' + sa_table.name.replace('"', '""') + '"'        # SQLite side

    count = 0
    raw = engine.raw_connection()
    try:
        cur = raw.cursor()
        with cur.copy(f"COPY {table_ident} ({col_list}) FROM STDIN") as copy:
            for row in src.execute(f"SELECT * FROM {src_ident}"):
                copy.write_row(tuple(row))
                count += 1
                if count % 100_000 == 0:
                    print(f"  {sa_table.name}: {count:,} rows...")
        raw.commit()
    finally:
        raw.close()
    return count


def main() -> None:
    dataset = get_settings().dataset
    sqlite_path = Path(f"seed/data/{dataset}.sqlite")

    if not sqlite_path.exists():
        sys.exit(
            f"{dataset} SQLite not found at {sqlite_path}.\n"
            f"Run `bash ctrl/seed/download.sh` on the host first."
        )

    engine = get_engine()
    src = sqlite3.connect(sqlite_path)
    src.row_factory = sqlite3.Row

    prep = engine.dialect.identifier_preparer
    schema_ident = prep.quote_schema(dataset)
    with engine.begin() as conn:
        conn.execute(text(f"DROP SCHEMA IF EXISTS {schema_ident} CASCADE"))
        conn.execute(text(f"CREATE SCHEMA {schema_ident}"))

    metadata = MetaData(schema=dataset)
    tables = [_discover_table(src, metadata, t, dataset) for t in _list_tables(src)]
    metadata.create_all(engine)
    print(f"created {len(tables)} tables in schema {dataset!r}")

    for sa_table in tables:
        count = _copy_table(engine, src, sa_table)
        print(f"  loaded {sa_table.name}: {count:,} rows")

    src.close()
    print("done.")


if __name__ == "__main__":
    main()