"""
migrate_to_pg.py
================
Runs ON the server via SSH with Python 3.12 virtualenv.
Uses psycopg v3 (psycopg package, not psycopg2).
1. Creates all PostgreSQL tables (col_<collection>)
2. Imports all JSON export files from mongo_export/
3. Creates indexes on frequently queried fields
"""
import os, sys, json, uuid

try:
    import psycopg  # psycopg v3
    PG_V = 3
except ImportError:
    import psycopg2 as psycopg
    PG_V = 2

PG_CONN_STR = "host=127.0.0.1 port=5432 dbname=dopebling_app user=dopebling_idgapp password=Romania95!"
EXPORT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "mongo_export")


def get_conn():
    return psycopg.connect(PG_CONN_STR, autocommit=False)


def table_name(col):
    return f"col_{col}"


def create_tables(conn):
    cur = conn.cursor()
    files = [f[:-5] for f in os.listdir(EXPORT_DIR) if f.endswith('.json')]
    for col in sorted(files):
        tbl = table_name(col)
        cur.execute(f"""
            CREATE TABLE IF NOT EXISTS {tbl} (
                id TEXT PRIMARY KEY,
                data JSONB NOT NULL,
                created_at TIMESTAMPTZ DEFAULT NOW()
            )
        """)
        print(f"  [TABLE] {tbl} ensured")
    conn.commit()
    cur.close()


def create_indexes(conn):
    cur = conn.cursor()
    indexes = [
        ("col_users", "username", "(data->>'username')"),
        ("col_users", "email", "(data->>'email')"),
        ("col_users", "role", "(data->>'role')"),
        ("col_users", "username_lower", "(data->>'username_lower')"),
        ("col_topics", "board_id", "(data->>'board_id')"),
        ("col_topics", "author_id", "(data->>'author_id')"),
        ("col_posts", "topic_id", "(data->>'topic_id')"),
        ("col_posts", "author_id", "(data->>'author_id')"),
        ("col_boards", "slug", "(data->>'slug')"),
        ("col_boards", "parent_id", "(data->>'parent_id')"),
        ("col_notifications", "user_id", "(data->>'user_id')"),
        ("col_dm_threads", "key", "(data->>'key')"),
        ("col_dms", "thread_id", "(data->>'thread_id')"),
        ("col_idg_clans", "name", "(data->>'name')"),
        ("col_idg_clan_members", "clan_id", "(data->>'clan_id')"),
        ("col_idg_clan_members", "user_id", "(data->>'user_id')"),
        ("col_idg_matches", "tournament_id", "(data->>'tournament_id')"),
        ("col_idg_matches", "status", "(data->>'status')"),
        ("col_idg_news", "slug", "(data->>'slug')"),
        ("col_idg_reborn_users", "username", "(data->>'username')"),
        ("col_points_log", "user_id", "(data->>'user_id')"),
        ("col_files", "owner_id", "(data->>'owner_id')"),
    ]
    for tbl, idx_name, expr in indexes:
        idx = f"idx_{tbl}_{idx_name}"
        try:
            cur.execute(f"CREATE INDEX IF NOT EXISTS {idx} ON {tbl} {expr}")
            print(f"  [INDEX] {idx}")
        except Exception as e:
            print(f"  [WARN] Index {idx}: {e}")
            conn.rollback()
            cur = conn.cursor()
    conn.commit()
    cur.close()


def import_collection(conn, col_name, docs):
    if not docs:
        return 0
    tbl = table_name(col_name)
    cur = conn.cursor()
    count = 0
    for doc in docs:
        doc_id = doc.get("id") or doc.get("_id")
        if not doc_id:
            doc_id = str(uuid.uuid4()).replace("-", "")
            doc["id"] = doc_id
        doc_id = str(doc_id)
        try:
            cur.execute(
                f"INSERT INTO {tbl} (id, data) VALUES (%s, %s) "
                f"ON CONFLICT (id) DO UPDATE SET data = EXCLUDED.data",
                (doc_id, json.dumps(doc))
            )
            count += 1
        except Exception as e:
            print(f"  [WARN] {tbl} id={doc_id[:20]}: {e}")
            conn.rollback()
            cur = conn.cursor()
    conn.commit()
    cur.close()
    return count


def main():
    print("=" * 60)
    print("MongoDB -> PostgreSQL Migration (psycopg v" + str(PG_V) + ")")
    print("=" * 60)
    print(f"Export dir: {EXPORT_DIR}")

    conn = get_conn()
    print(f"\n[OK] Connected to PostgreSQL: dopebling_app@127.0.0.1")

    print("\n--- Creating tables ---")
    create_tables(conn)

    print("\n--- Importing data ---")
    total = 0
    files = sorted([f for f in os.listdir(EXPORT_DIR) if f.endswith('.json')])
    for fname in files:
        col_name = fname[:-5]
        fpath = os.path.join(EXPORT_DIR, fname)
        with open(fpath, 'r', encoding='utf-8') as f:
            docs = json.load(f)
        count = import_collection(conn, col_name, docs)
        total += count
        if count > 0:
            print(f"  [OK] {col_name}: {count} docs")
        else:
            print(f"  [--] {col_name}: empty")

    print("\n--- Creating indexes ---")
    create_indexes(conn)

    conn.close()
    print(f"\n[DONE] Total: {total} documents imported!")
    print("=" * 60)


if __name__ == "__main__":
    main()