"""
BizNode — SQLite to PostgreSQL Migration Script
================================================
Reads all data from memory/biznode.db (SQLite) and writes it to a
PostgreSQL database. Safe to re-run: uses ON CONFLICT DO NOTHING.
SQLite data is NEVER deleted or modified by this script.

Usage:
    python scripts/migrate_to_postgresql.py

Prerequisites:
    pip install psycopg2-binary
    PostgreSQL running with database + user already created.

Required env vars (set in .env or export before running):
    DB_TYPE=postgresql
    DB_HOST=localhost
    DB_PORT=5432
    DB_NAME=biznode
    DB_USER=biznode_user
    DB_PASSWORD=your_password

Or pass flags directly:
    python scripts/migrate_to_postgresql.py \\
        --pg-host localhost --pg-port 5432 \\
        --pg-db biznode --pg-user biznode_user --pg-password your_password
"""

import os
import sys
import sqlite3
import argparse

# ── Project root ──────────────────────────────────────────────────────────────
_PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, _PROJECT_ROOT)

# ── Load .env before anything that reads env vars ─────────────────────────────
def _load_env():
    env_path = os.path.join(_PROJECT_ROOT, ".env")
    if os.path.exists(env_path):
        with open(env_path, encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if "=" in line and not line.startswith("#"):
                    k, v = line.split("=", 1)
                    os.environ.setdefault(k.strip(), v.strip())

_load_env()

# ── Tables — migration order (parents before children) ───────────────────────
# sqlite_sequence is SQLite-internal and always skipped.
SKIP_TABLES = {"sqlite_sequence"}

# Tables whose PRIMARY KEY is TEXT (not SERIAL) — no sequence reset needed.
TEXT_PK_TABLES = {"trusted_nodes", "blocked_customers"}

ALL_TABLES = [
    "agent_identity",
    "backup_config",
    "businesses",
    "associates",
    "notes",
    "note_links",
    "leads",
    "owner_actions",
    "network_interactions",
    "sync_registry",
    "audit_logs",
    "query_logs",
    "email_instructions",
    "conversation_history",
    "otp_codes",
    "scheduled_tasks",
    "task_permissions",
    "mcp_keys",
    "trusted_nodes",
    "blocked_customers",
    "workflows",
    "workflow_steps",
    "user_workflow_state",
    "menu_items",
    "backup_history",
    "services",
    "payment_accounts",
    "payment_requests",
]


# ── Bootstrap: load database.py via importlib to bypass memory/__init__.py ────
def _load_db_module():
    import importlib.util
    spec = importlib.util.spec_from_file_location(
        "_db", os.path.join(_PROJECT_ROOT, "memory", "database.py")
    )
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


# ── Helpers ───────────────────────────────────────────────────────────────────

def _bar(char="=", width=60):
    return char * width


def _pg_connect(host, port, dbname, user, password):
    try:
        import psycopg2
    except ImportError:
        print()
        print("[ERROR] psycopg2-binary is not installed.")
        print("        Run: pip install psycopg2-binary")
        sys.exit(1)
    try:
        conn = psycopg2.connect(
            host=host, port=port, dbname=dbname,
            user=user, password=password, connect_timeout=10,
        )
        conn.autocommit = False
        return conn
    except Exception as e:
        print(f"\n[ERROR] Cannot connect to PostgreSQL: {e}")
        print(f"        host={host} port={port} dbname={dbname} user={user}")
        sys.exit(1)


def _sqlite_connect():
    db_path = os.environ.get("SQLITE_PATH",
                             os.path.join(_PROJECT_ROOT, "memory", "biznode.db"))
    if not os.path.exists(db_path):
        print(f"\n[ERROR] SQLite database not found: {db_path}")
        sys.exit(1)
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    return conn, db_path


def _get_sqlite_tables(sqlite_conn):
    cur = sqlite_conn.cursor()
    cur.execute("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name")
    return {r[0] for r in cur.fetchall()} - SKIP_TABLES


def _reset_sequence(pg_conn, table):
    """Reset the PostgreSQL SERIAL sequence to MAX(id) so future inserts don't collide."""
    cur = pg_conn.cursor()
    try:
        cur.execute(
            f"SELECT setval(pg_get_serial_sequence('{table}', 'id'), "
            f"COALESCE(MAX(id), 1)) FROM {table}"
        )
        pg_conn.commit()
    except Exception:
        pg_conn.rollback()   # table may not have an 'id' SERIAL — that's fine
    finally:
        cur.close()


def _migrate_table(table, sqlite_conn, pg_conn, verbose=True):
    """
    Migrate one table from SQLite to PostgreSQL.
    Returns (rows_read, rows_inserted) tuple.
    """
    # ── Read from SQLite ──────────────────────────────────────────────────────
    sl_cur = sqlite_conn.cursor()
    try:
        sl_cur.execute(f'SELECT * FROM "{table}"')
    except sqlite3.OperationalError:
        return 0, 0   # table doesn't exist in this SQLite DB version

    rows = sl_cur.fetchall()
    if not rows:
        return 0, 0

    cols = [d[0] for d in sl_cur.description]

    # ── Build INSERT ──────────────────────────────────────────────────────────
    col_list     = ", ".join(f'"{c}"' for c in cols)
    placeholders = ", ".join(["%s"] * len(cols))
    insert_sql   = (
        f'INSERT INTO "{table}" ({col_list}) VALUES ({placeholders}) '
        f'ON CONFLICT DO NOTHING'
    )

    # ── Write to PostgreSQL ───────────────────────────────────────────────────
    pg_cur = pg_conn.cursor()
    inserted = 0
    try:
        for row in rows:
            values = tuple(row)
            pg_cur.execute(insert_sql, values)
            inserted += pg_cur.rowcount
        pg_conn.commit()
    except Exception as e:
        pg_conn.rollback()
        print(f"  [FAIL] {table}: {e}")
        return len(rows), 0
    finally:
        pg_cur.close()

    # Reset SERIAL sequence so future auto-inserts don't conflict
    if table not in TEXT_PK_TABLES and "id" in cols:
        _reset_sequence(pg_conn, table)

    return len(rows), inserted


def _verify(sqlite_conn, pg_conn, tables_done):
    """Print row-count comparison table between SQLite and PostgreSQL."""
    print()
    print("  Verification — row counts")
    print(f"  {'Table':<35} {'SQLite':>8} {'PostgreSQL':>10} {'Status':>8}")
    print(f"  {'-'*35} {'-'*8} {'-'*10} {'-'*8}")

    all_ok = True
    for table in tables_done:
        sl_cur = sqlite_conn.cursor()
        pg_cur = pg_conn.cursor()
        try:
            sl_cur.execute(f'SELECT COUNT(*) FROM "{table}"')
            sl_count = sl_cur.fetchone()[0]

            pg_cur.execute(f'SELECT COUNT(*) FROM "{table}"')
            pg_row = pg_cur.fetchone()
            pg_count = pg_row[0] if pg_row else 0

            status = "OK" if sl_count == pg_count else "MISMATCH"
            if status == "MISMATCH":
                all_ok = False
            print(f"  {table:<35} {sl_count:>8} {pg_count:>10} {status:>8}")
        except Exception as e:
            print(f"  {table:<35} {'?':>8} {'?':>10} {'ERR':>8}  ({e})")
            all_ok = False
        finally:
            sl_cur.close()
            pg_cur.close()

    return all_ok


# ── Main ──────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="Migrate BizNode SQLite data to PostgreSQL")
    parser.add_argument("--pg-host",     default=os.environ.get("DB_HOST", "localhost"))
    parser.add_argument("--pg-port",     default=int(os.environ.get("DB_PORT") or "5432"), type=int)
    parser.add_argument("--pg-db",       default=os.environ.get("DB_NAME", "biznode"))
    parser.add_argument("--pg-user",     default=os.environ.get("DB_USER", "postgres"))
    parser.add_argument("--pg-password", default=os.environ.get("DB_PASSWORD", ""))
    parser.add_argument("--dry-run",     action="store_true",
                        help="Show what would be migrated without writing anything")
    args = parser.parse_args()

    print()
    print(_bar())
    print("  BizNode — SQLite to PostgreSQL Migration")
    print(_bar())
    print()

    # ── Connect ───────────────────────────────────────────────────────────────
    sqlite_conn, db_path = _sqlite_connect()
    print(f"  [OK] SQLite   : {db_path}")

    if args.dry_run:
        print("  [DRY-RUN] No data will be written to PostgreSQL.")
    else:
        pg_conn = _pg_connect(args.pg_host, args.pg_port,
                              args.pg_db, args.pg_user, args.pg_password)
        print(f"  [OK] PostgreSQL: {args.pg_host}:{args.pg_port}/{args.pg_db}")

    # ── Create schema on PostgreSQL via init_db() ─────────────────────────────
    if not args.dry_run:
        print()
        print("  Creating schema on PostgreSQL (init_db)...")
        os.environ["DB_TYPE"]     = "postgresql"
        os.environ["DB_HOST"]     = args.pg_host
        os.environ["DB_PORT"]     = str(args.pg_port)
        os.environ["DB_NAME"]     = args.pg_db
        os.environ["DB_USER"]     = args.pg_user
        os.environ["DB_PASSWORD"] = args.pg_password
        try:
            _db = _load_db_module()
            _db.init_db()
            print("  [OK] Schema created (all tables exist)")
        except Exception as e:
            print(f"  [WARN] init_db had an issue: {e}")
            print("         Continuing — tables may already exist.")

    # ── Get tables in SQLite ──────────────────────────────────────────────────
    sqlite_tables = _get_sqlite_tables(sqlite_conn)
    migrate_list  = [t for t in ALL_TABLES if t in sqlite_tables]
    # Also catch any tables in SQLite not in our list (future tables)
    extra = sqlite_tables - set(migrate_list) - SKIP_TABLES
    if extra:
        migrate_list.extend(sorted(extra))

    print()
    print(f"  Migrating {len(migrate_list)} tables...")
    print()

    # ── Migrate each table ────────────────────────────────────────────────────
    total_rows = 0
    total_ins  = 0
    tables_with_data = []

    for table in migrate_list:
        if args.dry_run:
            sl_cur = sqlite_conn.cursor()
            try:
                sl_cur.execute(f'SELECT COUNT(*) FROM "{table}"')
                count = sl_cur.fetchone()[0]
            except Exception:
                count = 0
            if count > 0:
                print(f"  [DRY-RUN] {table}: {count} rows would be migrated")
                tables_with_data.append(table)
            else:
                print(f"  [SKIP]    {table}: 0 rows")
        else:
            read, ins = _migrate_table(table, sqlite_conn, pg_conn)
            if read == 0:
                print(f"  [SKIP]    {table}: 0 rows")
            else:
                print(f"  [OK]      {table}: {read} rows read, {ins} inserted")
                total_rows += read
                total_ins  += ins
                tables_with_data.append(table)

    # ── Verify ────────────────────────────────────────────────────────────────
    if not args.dry_run and tables_with_data:
        all_ok = _verify(sqlite_conn, pg_conn, tables_with_data)

    # ── Close ─────────────────────────────────────────────────────────────────
    sqlite_conn.close()
    if not args.dry_run:
        pg_conn.close()

    # ── Summary ───────────────────────────────────────────────────────────────
    print()
    print(_bar())
    if args.dry_run:
        print("  DRY-RUN complete. No data was written.")
    else:
        print(f"  Migration complete.")
        print(f"  Rows read from SQLite    : {total_rows}")
        print(f"  Rows inserted (new only) : {total_ins}")
        if total_rows > total_ins:
            print(f"  Rows skipped (existing)  : {total_rows - total_ins}")
        print()
        if all_ok:
            print("  [OK] All row counts match.")
        else:
            print("  [WARN] Some counts differ — check the table above.")
    print(_bar())
    print()
    print("  Next steps:")
    print("  1. Open Dashboard (http://localhost:7777) > Setup > Database")
    print("     Set DB_TYPE=postgresql and fill in host/port/name/user/password")
    print("     Click Save Settings")
    print("  2. Restart services:")
    print("       docker compose up -d --force-recreate")
    print("  3. Verify at http://localhost:7777 — dashboard loads with your data")
    print("  4. Keep memory/biznode.db safe until PostgreSQL is confirmed working")
    print("  5. To rollback: set DB_TYPE=sqlite in .env and restart")
    print()


if __name__ == "__main__":
    main()
