""" Raw SQL-based Postgres DB I/O utils """

import os
from datetime import datetime
from typing import NamedTuple

import pandas as pd
import psycopg
from psycopg.rows import namedtuple_row

import pylogg

log = pylogg.get('postgres')


class Tunnel(NamedTuple):
    """ Configuration for setting up a SSH tunnel. """
    host : str          = os.environ.get('TUNNEL_HOST', None)
    port : int          = int(os.environ.get('TUNNEL_PORT', 5432))
    user : str          = os.environ.get('TUNNEL_USER', '')
    pswd : str          = os.environ.get('TUNNEL_PASS', '')


class Server(NamedTuple):
    """ Configuration for connecting to a database. """
    host : str          = os.environ.get('PGDB_HOST', "localhost")
    port : int          = int(os.environ.get('PGDB_PORT', 5432))
    name : str          = os.environ.get('PGDB_NAME', "postgres")
    user : str          = os.environ.get('PGDB_USER', "admin")
    pswd : str          = os.environ.get('PGDB_PASS', "")


def ssh_tunnel(tunnel : Tunnel, ssh_port : int = 22) -> dict:
    """ Create a SSH tunnel to a remote host.
        Returns: Local tunnel endpoint dict[host, port].
    """

    from sshtunnel import SSHTunnelForwarder

    log.trace("Creating SSH tunnel to", tunnel.host)
    proxy = SSHTunnelForwarder((tunnel.host, int(ssh_port)),
            ssh_username=tunnel.user,
            ssh_password=tunnel.pswd,
            remote_bind_address=('127.0.0.1', tunnel.port),
        )
    proxy.start()
    log.done("SSH tunnel :{} ==> {}", proxy.local_bind_port, tunnel.host)
    return dict(host="127.0.0.1", port=proxy.local_bind_port)


def connect(db : Server) -> psycopg.Connection:
    """ Connect to a specific database.
        Returns: The db connection object.
    """
    log.trace("Connecting to", db.name)

    conninfo =  f'host={db.host} '\
                f'port={db.port} '\
                f'dbname={db.name} '\
                f'user={db.user} '\
                f'password={db.pswd}'

    conn = psycopg.connect(conninfo)

    log.done("Connect to database:", db.name)
    return conn


def raw_sql(conn : psycopg.Connection, psycopg_sql_string, **kwargs) -> list:
    """ Perform SQL select query using a psycopg3 style SQL statement.
        Returns the list of selected rows.
    """
    log.trace("Executing SQL:", psycopg_sql_string)

    conn.autocommit = True
    conn.row_factory = namedtuple_row
    cur : psycopg.Cursor = conn.cursor()
    rows = cur.execute(query=psycopg_sql_string, params=kwargs).fetchall()

    return rows


def query(conn : psycopg.Connection, sql, **kwargs) -> pd.DataFrame:
    """ Query the database using raw sql.
        Return a pandas dataframe containing the results.
    """
    results = raw_sql(conn, sql, **kwargs)
    if not results:
        return pd.DataFrame()
    return pd.DataFrame.from_records(results, columns=results[0]._fields)


class Table:
    def __init__(self, _tableName, **cols) -> None:
        """ Generic table definition for creating and inserting.
            Example:
                tabl = Table("test",
                    name = "varchar NOT NULL UNIQUE",
                    age  = "int4 NOT NULL"
                )
        """
        self._tblname = _tableName
        self._columns = {
            "id": f"int4 NOT NULL GENERATED BY DEFAULT AS IDENTITY(INCREMENT BY 1 MINVALUE 1 MAXVALUE 2147483647 START 1 CACHE 1 NO CYCLE)",
            "date_added": "timestamptz NOT NULL",
        }
        self._indices = {}
        self._constraints = {"id": f"{self._tblname}_pkey PRIMARY KEY"}

        for k, v in cols.items():
            self._columns[k] = v

    def __repr__(self) -> str:
        r = f"Table {self._tblname}({', '.join([k for k in self._columns.keys()])})"
        return r

    def constraint(self, col : str, value : str):
        self._constraints[col] = value

    def index(self, col : str, ops : str = None):
        if ops: assert ops in ['varchar', 'text']
        self._indices[col] = f"{ops}_pattern_ops" if ops else ""

    def _create_sql(self) -> str:
        sql = f"CREATE TABLE IF NOT EXISTS {self._tblname} (\n\t"
        rows = [f'"{k}" {v}' for k, v in self._columns.items()]
        rows += [f'CONSTRAINT {v} ("{k}")' for k, v in self._constraints.items()]
        sql += ",\n\t".join(rows) + "\n);"
        return sql

    def _index_sql(self) -> list[str]:
        return [
            f'CREATE INDEX IF NOT EXISTS ix_{self._tblname}_{k} ON {self._tblname} USING btree ("{k}" {v});'
            for k, v in self._indices.items()
        ]

    def create_all(self, conn : psycopg.Connection, drop_existing = False):
        """ Create the table and all the defined indices.
        """
        log.trace("Creating table", self._tblname)

        try:
            if drop_existing:
                log.warn("Dropping existing table:", self._tblname)

                for k in self._indices.keys():
                    sql = f"DROP INDEX IF EXISTS ix_{self._tblname}_{k};"
                    conn.execute(sql)

                sql = f"DROP TABLE IF EXISTS {self._tblname};"
                conn.execute(sql)

            for sql in [self._create_sql()] + self._index_sql():
                log.trace(sql)
                conn.execute(sql)

            conn.commit()
            log.done("Table create:", self._tblname)

        except:
            conn.rollback()
            log.error("Failed to create table:", self._tblname)
            raise

    def insert_row(self, conn : psycopg.Connection, **kwargs):
        """ Insert a single table row and commit immediately.
            Returns the inserted Row with the id field. Example:
            tabl.insert_row(name = "John", age = 31)
        """
        log.trace("Insert row into", self._tblname)

        if "date_added" not in kwargs:
            kwargs["date_added"] = datetime.now()

        columns = [k for k in kwargs.keys()]
        for col in columns:
            assert col in self._columns, f"Column '{col}' not in {self}"

        column_string = ", ".join(columns)
        values_string = ", ".join([f"%({k})s" for k in columns])

        sql = f"""
        INSERT INTO {self._tblname} ({column_string})
        VALUES ({values_string}) RETURNING id;
        """
        row = raw_sql(conn, sql, **kwargs)

        if row:
            return row[0]
        return None

    def copy(self, conn : psycopg.Connection, rows : list[dict]):
        """ Insert a batch of rows and commit at the end. Example:
            tabl.copy([
                {'name': "Mike", 'age': 23},
                {'name': "Kate", 'age': 41},
                {'name': "Terry", 'age': 9},
            ])
        """

        t1 = log.info("Copying {} rows into {}", len(rows), self._tblname)

        columns = [k for k in rows[0].keys()]
        if "date_added" not in columns:
            columns.append('date_added')

        for col in columns:
            assert col in self._columns, f"Column '{col}' not in {self}"

        column_string = ", ".join(columns)
        sql = f"COPY {self._tblname} ({column_string}) FROM STDIN"

        cur : psycopg.Cursor = conn.cursor()
        try:
            with cur.copy(sql) as copy:
                for row in rows:
                    if 'date_added' not in row:
                        row['date_added'] = datetime.now()
                    copy.write_row([row[k] for k in columns])

            conn.commit()
            t1.done("Copy {} rows to {}", len(rows), self._tblname)
        except:
            conn.rollback()
            log.error("Failed to copy rows to", self._tblname)
            raise


if __name__ == "__main__":
    server = Server()

    if os.environ.get("TUNNEL_HOST", None):
        tunnel = Tunnel()
        server = server._replace(**ssh_tunnel(tunnel))

    tabl = Table("test",
        name   = "varchar NOT NULL UNIQUE",
        age    = "int4 NOT NULL"
    )

    tabl.index('age')
    tabl.index('name', 'varchar')

    db = connect(server)
    tabl.create_all(db, drop_existing=True)

    print(tabl.insert_row(db, name = "John", age = 31))
    print(tabl.insert_row(db, name = "Doe", age = 32))

    tabl.copy(db, [
        {'name': "Mike", 'age': 23},
        {'name': "Kate", 'age': 41},
        {'name': "Terry", 'age': 9},
    ])

    for row in raw_sql(db, "SELECT * FROM test;"):
        print(row)

    print("Done")
