#!/usr/bin/env python3
"""Convert a MongoDB collection into a PostgreSQL table.

Contains `convert_mongo_to_pg(mongo_uri, mongo_db, collection_name, pg_dsn, pg_table=None, ...)`
and a small CLI.

Notes:
- Uses `pymongo` to read documents and `psycopg2` to write to Postgres.
- Infers basic column types from a sample of documents; nested objects/arrays
  are stored as `JSONB`.
"""

from __future__ import annotations

import argparse
import json
import logging
from typing import Any, Dict, Iterable, List, Optional

import pymongo
import psycopg2
import psycopg2.extras as pgextras
import psycopg2.sql as sql

LOG = logging.getLogger("mongo2pgsql")


def _detect_type(value: Any) -> str:
    if value is None:
        return "null"
    if isinstance(value, bool):
        return "boolean"
    if isinstance(value, int) and not isinstance(value, bool):
        return "integer"
    if isinstance(value, float):
        return "real"
    if isinstance(value, (dict, list)):
        return "jsonb"
    return "text"


def _merge_types(types: List[str]) -> str:
    # priority: jsonb > text > real > integer > boolean > null
    s = set(types)
    if "jsonb" in s:
        return "jsonb"
    if "text" in s:
        return "text"
    if "real" in s:
        return "real"
    if "integer" in s:
        return "integer"
    if "boolean" in s:
        return "boolean"
    return "text"


def _pg_type_for(kind: str) -> str:
    return {
        "integer": "BIGINT",
        "real": "DOUBLE PRECISION",
        "boolean": "BOOLEAN",
        "jsonb": "JSONB",
        "text": "TEXT",
    }.get(kind, "TEXT")


def convert_mongo_to_pg(
    mongo_uri: str,
    mongo_db: str,
    collection_name: str,
    pg_dsn: str,
    pg_table: Optional[str] = None,
    sample_size: int = 500,
    batch_size: int = 1000,
    create_table: bool = True,
    pk_field: str = "_id",
) -> None:
    """Copy a MongoDB collection to PostgreSQL.

    - `pg_table` defaults to `collection_name` if omitted.
    - Nested documents/arrays are stored as JSONB.
    - `_id` is stored as TEXT by default.
    """
    if pg_table is None:
        pg_table = collection_name

    client = pymongo.MongoClient(mongo_uri)
    coll = client[mongo_db][collection_name]

    sample = list(coll.find({}, projection=None, limit=sample_size))
    if not sample:
        raise ValueError("collection is empty or not found")

    # gather all field names and observed types
    field_types: Dict[str, List[str]] = {}
    for doc in sample:
        for k, v in doc.items():
            t = _detect_type(v)
            field_types.setdefault(k, []).append(t)

    # finalize types
    final_types: Dict[str, str] = {}
    for k, types in field_types.items():
        merged = _merge_types(types)
        final_types[k] = _pg_type_for(merged)

    # ensure primary key present
    if pk_field not in final_types:
        final_types[pk_field] = "TEXT"

    cols = list(final_types.items())

    # create table
    if create_table:
        # Build CREATE TABLE statement using psycopg2.sql for safe identifiers
        col_defs = []
        for name, ptype in cols:
            col_defs.append(sql.SQL("{} {}").format(sql.Identifier(name), sql.SQL(ptype)))

        create_body = sql.SQL(", ").join(col_defs)
        if pk_field in final_types:
            pk_clause = sql.SQL(", PRIMARY KEY ({})").format(sql.SQL(",").join(map(sql.Identifier, [pk_field])))
            create_body = create_body + pk_clause

        create_stmt = sql.SQL("CREATE TABLE IF NOT EXISTS {} ({})").format(sql.Identifier(pg_table), create_body)
        with psycopg2.connect(pg_dsn) as pgconn:
            with pgconn.cursor() as cur:
                LOG.info("Creating table %s", pg_table)
                cur.execute(create_stmt)
                pgconn.commit()

    # stream documents and insert
    insert_cols = [n for n, _ in cols]

    with psycopg2.connect(pg_dsn) as pgconn:
        with pgconn.cursor() as cur:
            # build INSERT statement for execute_values: it expects a single %s
            base_insert = sql.SQL("INSERT INTO {} ({}) VALUES %s").format(
                sql.Identifier(pg_table),
                sql.SQL(', ').join(map(sql.Identifier, insert_cols)),
            ).as_string(pgconn)
            batch = []
            # Use an explicit session when using no_cursor_timeout to avoid pymongo warning
            with client.start_session() as session:
                cursor = coll.find({}, no_cursor_timeout=True, session=session).batch_size(batch_size)
                try:
                    for doc in cursor:
                        row = []
                        for col_name in insert_cols:
                            v = doc.get(col_name)
                            if v is None:
                                row.append(None)
                                continue
                            # convert ObjectId and other non-JSON types
                            if isinstance(v, (dict, list)):
                                row.append(pgextras.Json(v))
                            else:
                                row.append(str(v) if col_name == pk_field else v)
                        batch.append(tuple(row))
                        if len(batch) >= batch_size:
                            pgextras.execute_values(cur, base_insert, batch, template=None, page_size=batch_size)
                            pgconn.commit()
                            batch = []
                finally:
                    try:
                        cursor.close()
                    except Exception:
                        pass
            if batch:
                pgextras.execute_values(cur, base_insert, batch, template=None, page_size=len(batch))
                pgconn.commit()

    LOG.info("Finished copying collection %s.%s to %s", mongo_db, collection_name, pg_table)


def _cli() -> None:
    parser = argparse.ArgumentParser(description="Copy a MongoDB collection to PostgreSQL")
    parser.add_argument("mongo_uri", help="MongoDB connection URI, e.g. mongodb://user:pw@host:27017")
    parser.add_argument("mongo_db", help="MongoDB database name")
    parser.add_argument("collection", help="Collection name to copy")
    parser.add_argument("pg_dsn", help="Postgres DSN, e.g. \"host=.. dbname=.. user=.. password=..\"")
    parser.add_argument("--pg-table", help="Target Postgres table name (defaults to collection name)")
    parser.add_argument("--sample", type=int, default=500, help="Sample size to infer schema")
    parser.add_argument("--batch", type=int, default=1000, help="Insert batch size")
    parser.add_argument("--no-create", dest="create", action="store_false", help="Don't attempt to create table in PG")
    parser.add_argument("--verbose", "-v", action="store_true")
    args = parser.parse_args()

    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO)

    convert_mongo_to_pg(
        args.mongo_uri,
        args.mongo_db,
        args.collection,
        args.pg_dsn,
        pg_table=args.pg_table,
        sample_size=args.sample,
        batch_size=args.batch,
        create_table=args.create,
    )


if __name__ == "__main__":
    _cli()
