"""Postgres helper for reading tables into pandas DataFrame.

Provides PgHandler class with methods for reading and basic DDL/DML.
"""
from typing import Optional, List, Any
import pandas as pd
import psycopg2
from psycopg2 import sql


class PgHandler:
    """Simple Postgres handler for reading tables.

    Parameters (pass to constructor):
      - host, port, user, password, database (db)
    """

    def __init__(self, *, host: str = "localhost", port: int = 5432, user: str = "postgres",
                 password: Optional[str] = None, database: Optional[str] = None):
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.database = database

        # internal runtime resources
        self._conn = None

    def _get_conn(self):
        """Lazily create and return a psycopg2 connection."""
        if self._conn is None:
            if self.database is None:
                raise ValueError("database must be provided")
            self._conn = psycopg2.connect(host=self.host, port=self.port, user=self.user,
                                          password=self.password, dbname=self.database)
        return self._conn

    def _get_admin_conn(self, dbname: str = 'postgres'):
        """Create a fresh connection to a given database name (useful for creating/dropping DBs).

        This returns a new connection instance (caller must close it).
        """
        return psycopg2.connect(host=self.host, port=self.port, user=self.user,
                                password=self.password, dbname=dbname)

    def close(self):
        """Close the underlying DB connection."""
        if self._conn is not None:
            try:
                self._conn.close()
            finally:
                self._conn = None

    def read_table(self, table: str, schema: str = "public", columns: Optional[List[str]] = None) -> pd.DataFrame:
        """Read a table (or specific columns) and return a DataFrame.
        """
        conn = self._get_conn()
        cols = "*" if not columns else ", ".join([f'"{c}"' for c in columns])
        q = f'SELECT {cols} FROM "{schema}"."{table}"'
        return pd.read_sql_query(q, con=conn)

    def read_table_by_conditions(self, table: str, conditions: Any, schema: str = "public", columns: Optional[List[str]] = None) -> pd.DataFrame:
        """Read rows filtered by condition dict or list of such dicts."""
        if isinstance(conditions, dict):
            condition_list = [conditions]
        elif isinstance(conditions, list):
            condition_list = conditions
        else:
            raise ValueError("conditions must be a dict or list of dicts")

        where_clauses = []
        params_list: List[Any] = []

        for cond in condition_list:
            col = cond.get('column')
            op = cond.get('condition')
            target = cond.get('target')
            if col is None or op is None or target is None:
                raise ValueError(f"condition dict missing keys: {cond}")

            if not isinstance(target, (list, tuple)) or len(target) < 2:
                raise ValueError(f"target must be a list like ['relative','col'] or ['direct', value]: {target}")

            kind = target[0]
            left = f'"{col}"'
            if kind == 'relative':
                other_col = target[1]
                right = f'"{other_col}"'
                if op == 'equal':
                    clause = f"{left} = {right}"
                elif op == 'larger than':
                    clause = f"{left} > {right}"
                elif op == 'lower than':
                    clause = f"{left} < {right}"
                else:
                    raise ValueError(f"unsupported condition operator: {op}")
            elif kind == 'direct':
                val = target[1]
                if op == 'equal':
                    clause = f"{left} = %s"
                elif op == 'larger than':
                    clause = f"{left} > %s"
                elif op == 'lower than':
                    clause = f"{left} < %s"
                else:
                    raise ValueError(f"unsupported condition operator: {op}")
                params_list.append(val)
            else:
                raise ValueError(f"target first element must be 'relative' or 'direct', got {kind}")

            where_clauses.append(clause)

        where_sql = " AND ".join(where_clauses) if where_clauses else ""
        cols = "*" if not columns else ", ".join([f'"{c}"' for c in columns])
        q = f'SELECT {cols} FROM "{schema}"."{table}"'
        if where_sql:
            q = q + " WHERE " + where_sql

        conn = self._get_conn()
        return pd.read_sql_query(q, con=conn, params=tuple(params_list) if params_list else None)

    def query(self, query: str) -> list[dict]:
        conn = self._get_conn()
        cursor = conn.cursor()
        cursor.execute(query)
        payload = cursor.fetchall()
        columns = [desc[0] for desc in cursor.description]
        result_list = [dict(zip(columns, row)) for row in payload]
        return result_list

    def get_all_tables(self, schema: str = "public") -> list[str]:
        conn = self._get_conn()
        cursor = conn.cursor()
        cursor.execute("select tablename from pg_tables where schemaname=%s", (schema,))
        tables = [row[0] for row in cursor.fetchall()]
        return tables

    # -------------------- DDL / DB operations --------------------
    def database_exists(self, db_name: str) -> bool:
        conn = self._get_admin_conn('postgres')
        try:
            cur = conn.cursor()
            cur.execute("SELECT 1 FROM pg_database WHERE datname = %s", (db_name,))
            return cur.fetchone() is not None
        finally:
            conn.close()

    def create_database(self, db_name: str, if_force_create: bool = False) -> bool:
        if self.database == db_name:
            pass
        if self.database_exists(db_name):
            if not if_force_create:
                return False
            self.drop_database(db_name)
        conn = self._get_admin_conn('postgres')
        conn.autocommit = True
        try:
            cur = conn.cursor()
            cur.execute(sql.SQL('CREATE DATABASE {}').format(sql.Identifier(db_name)))
            return True
        finally:
            conn.close()

    def drop_database(self, db_name: str) -> None:
        conn = self._get_admin_conn('postgres')
        conn.autocommit = True
        try:
            cur = conn.cursor()
            cur.execute("SELECT pid FROM pg_stat_activity WHERE datname = %s AND pid <> pg_backend_pid()", (db_name,))
            rows = cur.fetchall()
            for (pid,) in rows:
                try:
                    cur.execute("SELECT pg_terminate_backend(%s)", (pid,))
                except Exception:
                    pass
            cur.execute(sql.SQL('DROP DATABASE IF EXISTS {}').format(sql.Identifier(db_name)))
        finally:
            conn.close()

    def schema_exists(self, schema: str) -> bool:
        conn = self._get_conn()
        cur = conn.cursor()
        cur.execute("SELECT 1 FROM pg_namespace WHERE nspname = %s", (schema,))
        return cur.fetchone() is not None

    def create_schema(self, schema: str, if_force_create: bool = False) -> bool:
        conn = self._get_conn()
        cur = conn.cursor()
        if self.schema_exists(schema):
            if not if_force_create:
                return False
            else:
                self.drop_schema(schema)
        cur.execute(sql.SQL('CREATE SCHEMA {}').format(sql.Identifier(schema)))
        conn.commit()
        return True

    def drop_schema(self, schema: str) -> None:
        conn = self._get_conn()
        cur = conn.cursor()
        cur.execute(sql.SQL('DROP SCHEMA IF EXISTS {} CASCADE').format(sql.Identifier(schema)))
        conn.commit()

    def table_exists(self, table: str, schema: str = 'public') -> bool:
        conn = self._get_conn()
        cur = conn.cursor()
        cur.execute(
            "SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_schema = %s AND table_name = %s)",
            (schema, table)
        )
        row = cur.fetchone()
        return bool(row and row[0])

    # alias to match some older call sites
    def exists_table(self, table_name: str, schema: str = 'public') -> bool:
        return self.table_exists(table_name, schema=schema)

    def create_table(self, table: str, columns_def: str, schema: str = 'public', if_force_create: bool = False) -> bool:
        """Create a table with provided columns definition SQL (e.g. 'id serial PRIMARY KEY, name text')."""
        conn = self._get_conn()
        cur = conn.cursor()
        if self.table_exists(table, schema):
            if not if_force_create:
                return False
            else:
                self.drop_table(table, schema=schema)
        q = sql.SQL('CREATE TABLE {}.{} ({})').format(sql.Identifier(schema), sql.Identifier(table), sql.SQL(columns_def))
        cur.execute(q)
        conn.commit()
        return True

    def create_table_from_dict(self, table_structure: dict, table_name: str, force_create: bool = False, primary_key_list: Optional[list] = None, schema: str = 'public') -> None:
        """Create table from dict mapping column name -> sql type. Preserves primary keys if provided."""
        table_name_l = table_name.lower()
        # build column definitions
        cols = []
        for key, val in table_structure.items():
            colname = key.replace(' ', '_').replace('-', '_')
            coldef = f'"{colname}" {val}'
            if primary_key_list and key in primary_key_list:
                coldef += ' PRIMARY KEY'
            cols.append(coldef)
        columns_def = ', '.join(cols)
        # delegate to create_table
        self.create_table(table_name_l, columns_def, schema=schema, if_force_create=force_create)

    def drop_table(self, table: str, schema: str = 'public') -> None:
        conn = self._get_conn()
        cur = conn.cursor()
        cur.execute(sql.SQL('DROP TABLE IF EXISTS {}.{} CASCADE').format(sql.Identifier(schema), sql.Identifier(table)))
        conn.commit()

    def insert_into_table(self, table_name: str, data_dict_list: list, batch: int = 1000, schema: str = 'public') -> None:
        """Insert list of dicts into table in batches."""
        if not data_dict_list:
            return
        table_name_l = table_name.lower()
        if not self.table_exists(table_name_l, schema=schema):
            raise ValueError(f'table: {schema}.{table_name_l} not exists in database.')
        conn = self._get_conn()
        cur = conn.cursor()
        columns = [k.replace(' ', '_').replace('-', '_') for k in data_dict_list[0].keys()]
        col_names = ','.join([f'"{c}"' for c in columns])
        placeholders = ','.join(['%s'] * len(columns))
        sql_stmt = f'INSERT INTO "{schema}"."{table_name_l}" ({col_names}) VALUES ({placeholders})'
        target = [list(item.values()) for item in data_dict_list]
        for i in range(0, len(target), batch):
            cur.executemany(sql_stmt, target[i:i+batch])
            conn.commit()

    def insert_json_into_table(self, table: str, json_list: list, schema: str = "public", batch: int = 1000, unique_columns: Optional[str] = None) -> None:
        """批量将 json（字典）列表插入到指定 schema 下的 table 中，支持 batch 和 unique_columns（ON CONFLICT ...）。"""
        if not json_list:
            return
        conn = self._get_conn()
        cur = conn.cursor()
        table_name_l = table.lower()
        if not self.table_exists(table_name_l, schema=schema):
            raise ValueError(f'table: {schema}.{table_name_l} not exists in database.')
        columns = [k.replace(' ', '_').replace('-', '_') for k in json_list[0].keys()]
        col_names = ','.join([f'"{c}"' for c in columns])
        placeholders = ','.join(['%s'] * len(columns))
        conflict_sql = f' ON CONFLICT ({unique_columns}) DO NOTHING' if unique_columns else ''
        base_sql = f'INSERT INTO "{schema}"."{table_name_l}" ({col_names}) VALUES ({placeholders}){conflict_sql}'
        target = [list(item.values()) for item in json_list]
        for i in range(0, len(target), batch):
            try:
                cur.executemany(base_sql, target[i:i+batch])
            except Exception:
                # fall back to row-by-row to avoid losing all rows on partial error
                for row in target[i:i+batch]:
                    try:
                        cur.execute(base_sql, row)
                    except Exception:
                        pass
            conn.commit()

    def __enter__(self):
        self._get_conn()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
        return False


if __name__ == '__main__':
    from src.config.load_config import project_config

    host = project_config.get('db_host') or project_config.get('host')
    port = project_config.get('db_port') or project_config.get('port')
    user = project_config.get('db_user') or project_config.get('user')
    password = project_config.get('db_password') or project_config.get('password')
    database = project_config.get('db_name') or project_config.get('database')

    pg_handler = PgHandler(
        host=host, port=int(port or 5432), user=user, password=password or None, database=database,
    )

    with pg_handler as pg:
        xxx = pg.get_all_tables()
        print(xxx)
