#!/usr/bin/env python3
"""Simple utility to copy a single table from SQLite to PostgreSQL.

Usage example:
    python sqlite2pgsql /path/to/sqlite.db "host=... dbname=... user=... password=..." mytable

Implements `convert_table_sqlite_to_pg(sqlite_path, pg_dsn, table_name, create_table=True, batch_size=1000)`
which copies schema (basic types) and data in batches.
"""

from __future__ import annotations

import argparse
import logging
import sqlite3
from typing import List, Tuple

import psycopg2
from psycopg2 import sql

LOG = logging.getLogger("sqlite2pgsql")


def _map_sqlite_type_to_pg(sqlite_type: str) -> str:
    t = (sqlite_type or "").upper()
    if "INT" in t:
        return "INTEGER"
    if any(x in t for x in ("CHAR", "CLOB", "TEXT")):
        return "TEXT"
    if "BLOB" in t:
        return "BYTEA"
    if any(x in t for x in ("REAL", "FLOA", "DOUB")):
        return "REAL"
    if "NUM" in t or "DEC" in t:
        return "NUMERIC"
    # default fallback
    return "TEXT"


def _get_table_info_sqlite(conn: sqlite3.Connection, table: str) -> List[Tuple]:
    cur = conn.execute(f"PRAGMA table_info('{table}')")
    rows = cur.fetchall()
    return rows


def list_tables_sqlite(sqlite_path: str) -> List[str]:
    """Return a list of non-system table names from the SQLite database."""
    conn = sqlite3.connect(sqlite_path)
    try:
        cur = conn.execute("SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%' ORDER BY name;")
        return [r[0] for r in cur.fetchall()]
    finally:
        conn.close()


def convert_table_sqlite_to_pg(
    sqlite_path: str,
    pg_dsn: str,
    table_name: str,
    create_table: bool = True,
    batch_size: int = 1000,
) -> None:
    """Copy a table from a SQLite file to PostgreSQL.

    Arguments:
        sqlite_path: path to the .db file
        pg_dsn: libpq connection string for psycopg2
        table_name: table to copy
        create_table: if True, attempt to create the table in Postgres (basic mapping)
        batch_size: number of rows per insert batch
    """
    LOG.info("Opening SQLite DB %s", sqlite_path)
    sconn = sqlite3.connect(sqlite_path)
    sconn.row_factory = sqlite3.Row

    try:
        info = _get_table_info_sqlite(sconn, table_name)
        if not info:
            raise ValueError(f"Table {table_name} not found in {sqlite_path}")

        columns = []
        pk_cols = []
        for col in info:
            # PRAGMA table_info: cid, name, type, notnull, dflt_value, pk
            name = col[1]
            ctype = col[2]
            notnull = bool(col[3])
            dflt = col[4]
            pk = bool(col[5])
            pgtype = _map_sqlite_type_to_pg(ctype)
            columns.append((name, pgtype, notnull, dflt))
            if pk:
                pk_cols.append(name)

        if create_table:
            LOG.info("Connecting to Postgres to create table %s", table_name)
            with psycopg2.connect(pg_dsn) as pgconn:
                with pgconn.cursor() as cur:
                    col_defs = []
                    for name, pgtype, notnull, _dflt in columns:
                        part = sql.SQL("{} {}{}") .format(
                            sql.Identifier(name),
                            sql.SQL(pgtype),
                            sql.SQL(" NOT NULL") if notnull else sql.SQL(""),
                        )
                        col_defs.append(part)

                    if pk_cols:
                        pk_clause = sql.SQL(", PRIMARY KEY ({})").format(sql.SQL(",").join(map(sql.Identifier, pk_cols)))
                        create_body = sql.SQL(", ").join(col_defs) + pk_clause
                    else:
                        create_body = sql.SQL(", ").join(col_defs)

                    create_stmt = sql.SQL("CREATE TABLE IF NOT EXISTS {} ({})").format(
                        sql.Identifier(table_name), create_body
                    )
                    cur.execute(create_stmt)
                    pgconn.commit()

        # Now copy data
        LOG.info("Copying data for table %s", table_name)
        s_cur = sconn.cursor()
        s_cur.execute(f"SELECT * FROM '{table_name}'")

        cols = [d[0] for d in s_cur.description]
        placeholders = ",".join(["%s"] * len(cols))
        insert_sql = sql.SQL("INSERT INTO {} ({}) VALUES ({})").format(
            sql.Identifier(table_name),
            sql.SQL(",").join(map(sql.Identifier, cols)),
            sql.SQL(placeholders),
        )

        with psycopg2.connect(pg_dsn) as pgconn:
            with pgconn.cursor() as pgcur:
                batch = []
                for row in s_cur:
                    batch.append(tuple(row))
                    if len(batch) >= batch_size:
                        pgcur.executemany(insert_sql.as_string(pgconn), batch)
                        pgconn.commit()
                        batch = []
                if batch:
                    pgcur.executemany(insert_sql.as_string(pgconn), batch)
                    pgconn.commit()

    finally:
        sconn.close()


def _cli() -> None:
    parser = argparse.ArgumentParser(description="Copy one table from SQLite to PostgreSQL")
    parser.add_argument("sqlite_db", help="Path to SQLite database file")
    parser.add_argument("pg_dsn", help="Postgres DSN, e.g. \"host=.. dbname=.. user=.. password=..\"")
    parser.add_argument("table", nargs="?", help="Table name to copy (optional if --tables or --all used)")
    parser.add_argument("--tables", help="Comma-separated list of tables to copy")
    parser.add_argument("--all", action="store_true", help="Copy all non-system tables from the SQLite DB")
    parser.add_argument("--no-create", dest="create", action="store_false", help="Don't attempt to create table in PG")
    parser.add_argument("--list", dest="list_tables", action="store_true", help="List tables in the SQLite DB and exit")
    parser.add_argument("--batch", type=int, default=1000, help="Batch size for inserts")
    parser.add_argument("--verbose", "-v", action="store_true", help="Verbose logging")
    args = parser.parse_args()

    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO)
    if getattr(args, "list_tables", False):
        for t in list_tables_sqlite(args.sqlite_db):
            print(t)
        return

    # Determine tables to process
    table_names = []
    if args.all:
        table_names = list_tables_sqlite(args.sqlite_db)
    elif args.tables:
        table_names = [t.strip() for t in args.tables.split(",") if t.strip()]
    elif args.table:
        table_names = [args.table]
    else:
        parser.error("must specify a table, --tables, or --all unless --list is used")

    if not args.pg_dsn:
        parser.error("pg_dsn is required to copy tables to Postgres")

    # Process tables sequentially
    for tbl in table_names:
        try:
            LOG.info("Processing table %s", tbl)
            convert_table_sqlite_to_pg(args.sqlite_db, args.pg_dsn, tbl, create_table=args.create, batch_size=args.batch)
        except Exception as exc:
            LOG.exception("Failed to process table %s: %s", tbl, exc)


if __name__ == "__main__":
    _cli()
