#!/usr/bin/env python3
"""Clone inventory tables from one tenant DB to another.

Uses COPY with the columns common to both source and target, so minor schema
mismatches (e.g. missing latitude/longitude in branches) are handled.

Example:
    python scripts/clone_inventory.py \
        --source tenant_autopartes_estrada \
        --target tenant_originales_autopartes
"""

import argparse
import sys
import tempfile
import time
from contextlib import closing

import psycopg2


# (table, source_where_clause)
TABLES = [
    ("branches", "id <> 1"),
    ("inventory", None),
    ("inventory_stock", None),
    ("inventory_sku_aliases", None),
    ("inventory_stock_summary", None),
    ("inventory_vehicle_compat", None),
]

SEQUENCES = [
    ("branches_id_seq", "branches"),
    ("inventory_id_seq", "inventory"),
    ("inventory_stock_id_seq", "inventory_stock"),
    ("inventory_sku_aliases_id_seq", "inventory_sku_aliases"),
    ("inventory_vehicle_compat_id_seq", "inventory_vehicle_compat"),
]


def connect(db_name: str):
    return psycopg2.connect(host="localhost", user="postgres", dbname=db_name)


def get_columns(cur, table: str):
    cur.execute(
        """
        SELECT column_name
        FROM information_schema.columns
        WHERE table_name = %s
          AND table_schema = 'public'
        ORDER BY ordinal_position
        """,
        (table,),
    )
    return [r[0] for r in cur.fetchall()]


def copy_table(src_conn, dst_conn, table: str, where: str | None):
    with src_conn.cursor() as src_cur, dst_conn.cursor() as dst_cur:
        src_cols = get_columns(src_cur, table)
        dst_cols = set(get_columns(dst_cur, table))
        common = [c for c in src_cols if c in dst_cols]
        if not common:
            print(f"Skipping {table}: no common columns")
            return

        col_sql = ", ".join(f'"{c}"' for c in common)
        copy_to = f'COPY (SELECT {col_sql} FROM "{table}"'
        if where:
            copy_to += f" WHERE {where}"
        copy_to += ") TO STDOUT"
        copy_from = f'COPY "{table}" ({col_sql}) FROM STDIN'

        print(f"Copying {table} ({len(common)} columns)...", end=" ", flush=True)
        start = time.time()

        with tempfile.SpooledTemporaryFile(max_size=50 * 1024 * 1024, mode="w+b") as tmp:
            src_cur.copy_expert(copy_to, tmp)
            tmp.seek(0)
            dst_cur.copy_expert(copy_from, tmp)

        dst_conn.commit()
        elapsed = time.time() - start
        print(f"done in {elapsed:.1f}s")


def reset_sequences(dst_conn):
    with dst_conn.cursor() as cur:
        for seq, table in SEQUENCES:
            cur.execute(
                f"SELECT setval('{seq}', COALESCE((SELECT MAX(id) FROM \"{table}\"), 1), true)"
            )
    dst_conn.commit()


def main():
    parser = argparse.ArgumentParser(description="Clone inventory between tenant DBs")
    parser.add_argument("--source", required=True)
    parser.add_argument("--target", required=True)
    args = parser.parse_args()

    src = connect(args.source)
    dst = connect(args.target)

    try:
        # Prepare target: remove extra branches, truncate inventory tables.
        with dst.cursor() as cur:
            print("Preparing target tables...", end=" ", flush=True)
            cur.execute("DELETE FROM branches WHERE id <> 1")
            cur.execute(
                """
                TRUNCATE TABLE inventory,
                              inventory_stock,
                              inventory_sku_aliases,
                              inventory_stock_summary,
                              inventory_vehicle_compat
                CASCADE
                """
            )
            dst.commit()
            print("done")

        for table, where in TABLES:
            copy_table(src, dst, table, where)

        print("Resetting sequences...", end=" ", flush=True)
        reset_sequences(dst)
        print("done")

        print("Inventory clone completed.")
    except Exception as e:
        print(f"ERROR: {e}", file=sys.stderr)
        dst.rollback()
        raise
    finally:
        src.close()
        dst.close()


if __name__ == "__main__":
    main()