#!/usr/bin/env python3
"""Utility to expose a PostgreSQL table to SQLite by copying it into a sqlite file.

This is a pragmatic "plugin": it creates a local SQLite table (read-only by default)
that mirrors a PostgreSQL table so SQLite queries can run locally against the data.

Usage examples:
  # copy single table into sqlite file
  python sqlite_pg_plugin.py copy --pg "host=... dbname=... user=... password=..." public.mytable /tmp/my.db

  # list tables in Postgres
  python sqlite_pg_plugin.py list --pg "host=... dbname=... user=... password=..."

This avoids needing to build a native virtual-table extension. For a direct virtual
table (no copy) a C extension using SQLite virtual table API + libpq would be needed.
"""

from __future__ import annotations

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

import psycopg2
from psycopg2 import sql

LOG = logging.getLogger("sqlite_pg_plugin")


def list_tables_pg(pg_dsn: str) -> List[str]:
    q = "SELECT table_schema||'.'||table_name FROM information_schema.tables WHERE table_type='BASE TABLE' AND table_schema NOT IN ('pg_catalog','information_schema') ORDER BY 1;"
    with psycopg2.connect(pg_dsn) as conn:
        with conn.cursor() as cur:
            cur.execute(q)
            return [r[0] for r in cur.fetchall()]


def fetch_table_schema_pg(pg_dsn: str, qualified_table: str) -> List[Tuple[str, str, bool]]:
    """Return list of (colname, pgtype, notnull) for the given qualified table (schema.table)."""
    if '.' in qualified_table:
        schema, table = qualified_table.split('.', 1)
    else:
        schema, table = 'public', qualified_table

    q = sql.SQL("SELECT column_name, data_type, is_nullable FROM information_schema.columns WHERE table_schema=%s AND table_name=%s ORDER BY ordinal_position;")
    with psycopg2.connect(pg_dsn) as conn:
        with conn.cursor() as cur:
            cur.execute(q, (schema, table))
            out = []
            for name, dtype, is_nullable in cur.fetchall():
                out.append((name, dtype, is_nullable == 'NO'))
            return out


def _map_pg_to_sqlite(pg_type: str) -> str:
    t = (pg_type or '').lower()
    if 'int' in t or t in ('bigserial', 'serial'):
        return 'INTEGER'
    if any(x in t for x in ('char', 'text', 'uuid')):
        return 'TEXT'
    if any(x in t for x in ('timestamp', 'date', 'time')):
        return 'TEXT'
    if any(x in t for x in ('numeric', 'decimal', 'real', 'double', 'float')):
        return 'REAL'
    if 'bool' in t:
        return 'INTEGER'
    return 'TEXT'


def copy_table_pg_to_sqlite(pg_dsn: str, qualified_table: str, sqlite_path: str, sqlite_table: str | None = None, batch_size: int = 1000, overwrite: bool = True) -> None:
    """Copy schema and data from Postgres table into a SQLite file.

    - `qualified_table` can be 'schema.table' or 'table'
    - `sqlite_table` is optional name to use in sqlite; defaults to table part
    """
    if '.' in qualified_table and not sqlite_table:
        sqlite_table = qualified_table.split('.', 1)[1]
    elif not sqlite_table:
        sqlite_table = qualified_table

    LOG.info('Fetching schema for %s', qualified_table)
    cols = fetch_table_schema_pg(pg_dsn, qualified_table)
    if not cols:
        raise ValueError(f'No columns found for {qualified_table}')

    # create sqlite table
    sconn = sqlite3.connect(sqlite_path)
    try:
        s_cur = sconn.cursor()
        if overwrite:
            s_cur.execute(f"DROP TABLE IF EXISTS '{sqlite_table}'")

        col_defs = []
        for name, pgtype, notnull in cols:
            stype = _map_pg_to_sqlite(pgtype)
            col_defs.append(f"{name} {stype}{' NOT NULL' if notnull else ''}")

        create_sql = f"CREATE TABLE IF NOT EXISTS '{sqlite_table}' ({', '.join(col_defs)})"
        s_cur.execute(create_sql)
        sconn.commit()

        # stream rows from Postgres
        with psycopg2.connect(pg_dsn) as pgconn:
            with pgconn.cursor(name='stream_cursor') as pgcur:
                pgcur.itersize = batch_size
                pgcur.execute(sql.SQL('SELECT {} FROM {}').format(sql.SQL(',').join(sql.Identifier(c[0]) for c in cols), sql.Identifier(*qualified_table.split('.', 1)) if '.' in qualified_table else sql.Identifier(qualified_table)))

                placeholders = ','.join('?' for _ in cols)
                insert_sql = f"INSERT INTO '{sqlite_table}' ({', '.join(n for n,_,_ in cols)}) VALUES ({placeholders})"
                batch = []
                for row in pgcur:
                    batch.append(tuple(row))
                    if len(batch) >= batch_size:
                        s_cur.executemany(insert_sql, batch)
                        sconn.commit()
                        batch = []
                if batch:
                    s_cur.executemany(insert_sql, batch)
                    sconn.commit()
    finally:
        sconn.close()


def _cli() -> None:
    parser = argparse.ArgumentParser(description='Expose Postgres table(s) to SQLite by copying them')
    sub = parser.add_subparsers(dest='cmd', required=True)

    p_list = sub.add_parser('list', help='List tables in Postgres')
    p_list.add_argument('--pg', required=True, help='Postgres DSN')

    p_copy = sub.add_parser('copy', help='Copy a Postgres table into sqlite file')
    p_copy.add_argument('--pg', required=True, help='Postgres DSN')
    p_copy.add_argument('table', help='Qualified table name (schema.table or table)')
    p_copy.add_argument('sqlite', help='Path to sqlite file to write to')
    p_copy.add_argument('--target', help='Target sqlite table name (defaults to original table)')
    p_copy.add_argument('--batch', type=int, default=1000)
    p_copy.add_argument('--no-overwrite', dest='overwrite', action='store_false')
    p_copy.add_argument('--verbose', '-v', action='store_true')

    args = parser.parse_args()
    logging.basicConfig(level=logging.DEBUG if getattr(args, 'verbose', False) else logging.INFO)

    if args.cmd == 'list':
        for t in list_tables_pg(args.pg):
            print(t)
        return

    if args.cmd == 'copy':
        copy_table_pg_to_sqlite(args.pg, args.table, args.sqlite, sqlite_table=args.target, batch_size=args.batch, overwrite=args.overwrite)


if __name__ == '__main__':
    _cli()