import json
import logging
from typing import Any, Dict, Iterable, List, Optional, Union
import pandas as pd
import vertica_python
from mnemosynecore.vault.client import (
    get_secret,
    get_secret_test,
    get_connection_as_json as _get_connection_as_json,
)
from mnemosynecore.vault.univ_conn import un_conn, un_conn_test
from os import listdir, path
HAS_AIRFLOW = False
try:
    from airflow.operators.python import get_current_context
    from airflow import DAG
    from airflow.utils.task_group import TaskGroup
    from airflow.operators.dummy import DummyOperator
    HAS_AIRFLOW = True
except ImportError:
    pass


def read_sql_file(file_path: str):
    if not path.exists(file_path):
        print('Error: no file ' + file_path)
        return None
    with open(file_path, 'r', encoding='utf-8') as f:
        return f.read()


def load_sql_tasks_from_dir(dir_sql: str, vertica_conn_id: str):
    if not HAS_AIRFLOW:
        raise ImportError("Для использования load_sql_tasks_from_dir необходимо установить Airflow")

    from airflow.providers.vertica.operators.vertica import VerticaOperator

    context = get_current_context()
    dag = context['dag']

    tasks = {}
    files_list = listdir(dir_sql)
    for filename in files_list:
        if not filename.endswith(".sql"):
            continue
        path_file = path.join(dir_sql, filename)
        sql = read_sql_file(path_file)
        if sql:
            sql_statements = sql.split(';')
            task_name = 'task_' + filename.replace('.sql', '_vertica')
            tasks[task_name] = VerticaOperator(
                sql=sql_statements,
                task_id=task_name,
                vertica_conn_id=vertica_conn_id,
                dag=dag,
            )
    return tasks


def vertica_dedupe(
        table_name: str,
        unique_keys: Union[str, List[str]],
        conn: Optional[vertica_python.Connection] = None,
        conn_id: Optional[str] = None,
        date_col: Optional[str] = None,
        keep: str = "last",
        commit: bool = True
):

    if isinstance(unique_keys, str):
        unique_keys = [unique_keys]

    close_conn = False
    if conn is None:
        if not conn_id:
            raise ValueError("Нужно указать conn или conn_id")
        conn = vertica_conn(conn_id)
        close_conn = True

    order_clause = f"ORDER BY {date_col} DESC" if date_col and keep == "last" else ""
    unique_cols = ", ".join(unique_keys)

    join_conditions = " AND ".join([f"t.{k} = x.{k}" for k in unique_keys])

    dedupe_sql = f"""
        DELETE FROM {table_name} t
        USING (
            SELECT {unique_cols}, ROW_NUMBER() OVER (PARTITION BY {unique_cols} {order_clause}) AS rn
            FROM {table_name}
        ) x
        WHERE {join_conditions}
          AND x.rn > 1
    """

    vertica_sql(conn=conn, sql=dedupe_sql, commit=commit)

    if close_conn:
        conn.close()

    print(f"Удаление дубликатов из {table_name} завершено по ключам: {unique_keys}")


def get_connection_as_json(conn_name):
    return _get_connection_as_json(conn_name)


def vertica_upsert(
        df: pd.DataFrame,
        table_name: str,
        unique_keys: Union[str, List[str]],
        conn=None,
        conn_id: Optional[str] = None,
        date_col: Optional[str] = None,
        days_back: Optional[int] = None,
        commit: bool = True
):
    if df.empty:
        print("Нет данных для обновления.")
        return

    if isinstance(unique_keys, str):
        unique_keys = [unique_keys]

    close_conn = False
    if conn is None:
        if not conn_id:
            raise ValueError("conn или conn_id обязателен")
        conn = vertica_conn(conn_id)
        close_conn = True

    temp_table = f"{table_name}_tmp"
    create_temp_sql = f"""
        DROP TABLE IF EXISTS {temp_table};
        CREATE LOCAL TEMP TABLE {temp_table} AS
        SELECT * FROM {table_name} WHERE 1=0;
    """
    vertica_sql(conn=conn, sql=create_temp_sql, commit=False)

    rows = [tuple(x) for x in df.to_numpy()]
    cols = ", ".join(df.columns)
    placeholders = ", ".join(["%s"] * len(df.columns))

    insert_sql = f"INSERT INTO {temp_table} ({cols}) VALUES ({placeholders})"
    vertica_sql(conn=conn, sql=insert_sql, params=rows, many=True, commit=False)

    if date_col and days_back:
        delete_sql = f"""
            DELETE FROM {table_name}
            WHERE {date_col} >= CURRENT_DATE - INTERVAL '{days_back} day';
        """
        vertica_sql(conn=conn, sql=delete_sql, commit=False)

    merge_conditions = " AND ".join([f"t.{k} = s.{k}" for k in unique_keys])
    update_assignments = ", ".join([f"{c} = s.{c}" for c in df.columns if c not in unique_keys])

    merge_sql = f"""
        MERGE INTO {table_name} t
        USING {temp_table} s
        ON {merge_conditions}
        WHEN MATCHED THEN UPDATE SET {update_assignments}
        WHEN NOT MATCHED THEN INSERT ({cols}) VALUES ({cols});
    """
    vertica_sql(conn=conn, sql=merge_sql, commit=False)

    if commit:
        conn.commit()

    if close_conn:
        conn.close()

    print(f"Обновление таблицы {table_name} завершено. {len(df)} строк обработано.")


def _build_vertica_connect_kwargs(cfg: Dict[str, Any]) -> Dict[str, Any]:
    extra = cfg.get("extra") or {}
    if isinstance(extra, str):
        try:
            extra = json.loads(extra)
        except Exception:
            extra = {}

    login = cfg.get("user") or cfg.get("login")
    kwargs = {
        "host": cfg["host"],
        "port": int(cfg.get("port", 5433)),
        "user": login,
        "password": cfg.get("password"),
        "database": cfg.get("schema") or cfg.get("database"),
    }
    kwargs.update(extra if isinstance(extra, dict) else {})
    return kwargs


def vertica_conn(conn_id: str):
    cfg = get_secret(conn_id)
    kwargs = _build_vertica_connect_kwargs(cfg)
    return vertica_python.connect(**kwargs)


def vertica_conn_test(
    conn_id: str,
    dir_path: Optional[str] = None,
    creds_path: Optional[str] = None,
):
    cfg = get_secret_test(conn_id, dir_path=dir_path, creds_path=creds_path)
    kwargs = _build_vertica_connect_kwargs(cfg)
    return vertica_python.connect(**kwargs)


def get_vertica_engine(conn_id: str):
    return un_conn(conn_id, "vertica")


def get_vertica_engine_test(
    conn_id: str,
    dir_path: Optional[str] = None,
    creds_path: Optional[str] = None,
):
    return un_conn_test(
        conn_id=conn_id,
        conn_type="vertica",
        dir_path=dir_path,
        creds_path=creds_path,
    )


def vertica_sql(
    *,
    conn_id: str = None,
    conn: vertica_python.Connection = None,
    sql: str,
    params=None,
    many: bool = False,
    commit: bool = True
):
    if not conn and not conn_id:
        raise ValueError("Нужно указать conn или conn_id")

    close_conn = False
    if not conn:
        conn = vertica_conn(conn_id)
        close_conn = True

    try:
        with conn.cursor() as cur:
            if many:
                cur.executemany(sql, params)
            else:
                cur.execute(sql, params)

        if commit:
            conn.commit()

    except Exception:
        conn.rollback()
        logging.exception("Ошибка выполнения SQL в Vertica")
        raise

    finally:
        if close_conn:
            conn.close()


def vertica_select(*, conn_id: str = None, conn=None, sql: str, params=None):
    if not conn and not conn_id:
        raise ValueError("Нужно указать conn или conn_id")

    close_conn = False
    if not conn:
        conn = vertica_conn(conn_id)
        close_conn = True

    try:
        with conn.cursor() as cur:
            cur.execute(sql, params)
            columns = [desc[0] for desc in cur.description]
            rows = cur.fetchall()
        import pandas as pd
        return pd.DataFrame(rows, columns=columns)

    finally:
        if close_conn:
            conn.close()


def split_sql_statements(sql_text: str) -> List[str]:
    if not sql_text:
        return []
    return [q.strip() for q in sql_text.split(";") if q.strip()]


def vertica_sql_file(
    *,
    file_path: str,
    conn_id: str = None,
    conn: vertica_python.Connection = None,
    params=None,
    commit: bool = True,
) -> int:
    sql_text = read_sql_file(file_path)
    if sql_text is None:
        raise FileNotFoundError(file_path)

    statements = split_sql_statements(sql_text)
    if not statements:
        return 0

    close_conn = False
    if conn is None:
        if not conn_id:
            raise ValueError("Нужно указать conn или conn_id")
        conn = vertica_conn(conn_id)
        close_conn = True

    try:
        for statement in statements:
            vertica_sql(conn=conn, sql=statement, params=params, commit=False)
        if commit:
            conn.commit()
    finally:
        if close_conn:
            conn.close()
    return len(statements)


def vertica_sql_dir(
    *,
    dir_sql: str,
    conn_id: str = None,
    conn: vertica_python.Connection = None,
    recursive: bool = False,
    commit: bool = True,
) -> Dict[str, int]:
    if not path.isdir(dir_sql):
        raise NotADirectoryError(dir_sql)

    sql_files: List[str] = []
    if recursive:
        for root, _, files in __import__("os").walk(dir_sql):
            for filename in files:
                if filename.endswith(".sql"):
                    sql_files.append(path.join(root, filename))
    else:
        for filename in listdir(dir_sql):
            if filename.endswith(".sql"):
                sql_files.append(path.join(dir_sql, filename))

    sql_files.sort()

    close_conn = False
    if conn is None:
        if not conn_id:
            raise ValueError("Нужно указать conn или conn_id")
        conn = vertica_conn(conn_id)
        close_conn = True

    result: Dict[str, int] = {}
    try:
        for sql_file in sql_files:
            executed = vertica_sql_file(
                file_path=sql_file,
                conn=conn,
                commit=False,
            )
            result[sql_file] = executed
        if commit:
            conn.commit()
    finally:
        if close_conn:
            conn.close()

    return result


def vertica_select_scalar(
    *,
    conn_id: str = None,
    conn: vertica_python.Connection = None,
    sql: str,
    params=None,
    default=None,
):
    if not conn and not conn_id:
        raise ValueError("Нужно указать conn или conn_id")

    close_conn = False
    if not conn:
        conn = vertica_conn(conn_id)
        close_conn = True

    try:
        with conn.cursor() as cur:
            cur.execute(sql, params)
            row = cur.fetchone()
            if not row:
                return default
            return row[0]
    finally:
        if close_conn:
            conn.close()


def vertica_insert_dataframe(
    *,
    df: pd.DataFrame,
    table_name: str,
    conn_id: str = None,
    conn: vertica_python.Connection = None,
    commit: bool = True,
    chunk_size: int = 5000,
) -> int:
    if not isinstance(df, pd.DataFrame):
        raise TypeError("df должен быть pandas.DataFrame")
    if df.empty:
        return 0
    if chunk_size < 1:
        raise ValueError("chunk_size должен быть >= 1")

    close_conn = False
    if conn is None:
        if not conn_id:
            raise ValueError("Нужно указать conn или conn_id")
        conn = vertica_conn(conn_id)
        close_conn = True

    cols = ", ".join(df.columns)
    placeholders = ", ".join(["%s"] * len(df.columns))
    insert_sql = f"INSERT INTO {table_name} ({cols}) VALUES ({placeholders})"
    rows = [tuple(x) for x in df.to_numpy()]

    try:
        for i in range(0, len(rows), chunk_size):
            chunk = rows[i:i + chunk_size]
            vertica_sql(conn=conn, sql=insert_sql, params=chunk, many=True, commit=False)
        if commit:
            conn.commit()
    finally:
        if close_conn:
            conn.close()

    return len(rows)


def vertica_table_exists(
    *,
    table_name: str,
    schema_name: Optional[str] = None,
    conn_id: str = None,
    conn: vertica_python.Connection = None,
) -> bool:
    if not conn and not conn_id:
        raise ValueError("Нужно указать conn или conn_id")

    close_conn = False
    if not conn:
        conn = vertica_conn(conn_id)
        close_conn = True

    try:
        with conn.cursor() as cur:
            if schema_name:
                cur.execute(
                    """
                    SELECT 1
                    FROM v_catalog.tables
                    WHERE table_schema = %s AND table_name = %s
                    LIMIT 1
                    """,
                    (schema_name, table_name),
                )
            else:
                cur.execute(
                    """
                    SELECT 1
                    FROM v_catalog.tables
                    WHERE table_name = %s
                    LIMIT 1
                    """,
                    (table_name,),
                )
            return cur.fetchone() is not None
    finally:
        if close_conn:
            conn.close()


def _resolve_vertica_test_conn(
    *,
    conn=None,
    conn_id: Optional[str] = None,
    dir_path: Optional[str] = None,
    creds_path: Optional[str] = None,
):
    if conn is not None:
        return conn, False
    if not conn_id:
        raise ValueError("Нужно указать conn или conn_id")
    return (
        vertica_conn_test(
            conn_id=conn_id,
            dir_path=dir_path,
            creds_path=creds_path,
        ),
        True,
    )


def vertica_sql_test(
    *,
    conn_id: str = None,
    conn: vertica_python.Connection = None,
    sql: str,
    params=None,
    many: bool = False,
    commit: bool = True,
    dir_path: Optional[str] = None,
    creds_path: Optional[str] = None,
):
    conn, close_conn = _resolve_vertica_test_conn(
        conn=conn,
        conn_id=conn_id,
        dir_path=dir_path,
        creds_path=creds_path,
    )
    try:
        return vertica_sql(
            conn=conn,
            sql=sql,
            params=params,
            many=many,
            commit=commit,
        )
    finally:
        if close_conn:
            conn.close()


def vertica_select_test(
    *,
    conn_id: str = None,
    conn=None,
    sql: str,
    params=None,
    dir_path: Optional[str] = None,
    creds_path: Optional[str] = None,
):
    conn, close_conn = _resolve_vertica_test_conn(
        conn=conn,
        conn_id=conn_id,
        dir_path=dir_path,
        creds_path=creds_path,
    )
    try:
        return vertica_select(conn=conn, sql=sql, params=params)
    finally:
        if close_conn:
            conn.close()


def vertica_select_scalar_test(
    *,
    conn_id: str = None,
    conn: vertica_python.Connection = None,
    sql: str,
    params=None,
    default=None,
    dir_path: Optional[str] = None,
    creds_path: Optional[str] = None,
):
    conn, close_conn = _resolve_vertica_test_conn(
        conn=conn,
        conn_id=conn_id,
        dir_path=dir_path,
        creds_path=creds_path,
    )
    try:
        return vertica_select_scalar(
            conn=conn,
            sql=sql,
            params=params,
            default=default,
        )
    finally:
        if close_conn:
            conn.close()


def vertica_sql_file_test(
    *,
    file_path: str,
    conn_id: str = None,
    conn: vertica_python.Connection = None,
    params=None,
    commit: bool = True,
    dir_path: Optional[str] = None,
    creds_path: Optional[str] = None,
) -> int:
    conn, close_conn = _resolve_vertica_test_conn(
        conn=conn,
        conn_id=conn_id,
        dir_path=dir_path,
        creds_path=creds_path,
    )
    try:
        return vertica_sql_file(
            file_path=file_path,
            conn=conn,
            params=params,
            commit=commit,
        )
    finally:
        if close_conn:
            conn.close()


def vertica_sql_dir_test(
    *,
    dir_sql: str,
    conn_id: str = None,
    conn: vertica_python.Connection = None,
    recursive: bool = False,
    commit: bool = True,
    dir_path: Optional[str] = None,
    creds_path: Optional[str] = None,
) -> Dict[str, int]:
    conn, close_conn = _resolve_vertica_test_conn(
        conn=conn,
        conn_id=conn_id,
        dir_path=dir_path,
        creds_path=creds_path,
    )
    try:
        return vertica_sql_dir(
            dir_sql=dir_sql,
            conn=conn,
            recursive=recursive,
            commit=commit,
        )
    finally:
        if close_conn:
            conn.close()


def vertica_insert_dataframe_test(
    *,
    df: pd.DataFrame,
    table_name: str,
    conn_id: str = None,
    conn: vertica_python.Connection = None,
    commit: bool = True,
    chunk_size: int = 5000,
    dir_path: Optional[str] = None,
    creds_path: Optional[str] = None,
) -> int:
    conn, close_conn = _resolve_vertica_test_conn(
        conn=conn,
        conn_id=conn_id,
        dir_path=dir_path,
        creds_path=creds_path,
    )
    try:
        return vertica_insert_dataframe(
            df=df,
            table_name=table_name,
            conn=conn,
            commit=commit,
            chunk_size=chunk_size,
        )
    finally:
        if close_conn:
            conn.close()


def vertica_table_exists_test(
    *,
    table_name: str,
    schema_name: Optional[str] = None,
    conn_id: str = None,
    conn: vertica_python.Connection = None,
    dir_path: Optional[str] = None,
    creds_path: Optional[str] = None,
) -> bool:
    conn, close_conn = _resolve_vertica_test_conn(
        conn=conn,
        conn_id=conn_id,
        dir_path=dir_path,
        creds_path=creds_path,
    )
    try:
        return vertica_table_exists(
            table_name=table_name,
            schema_name=schema_name,
            conn=conn,
        )
    finally:
        if close_conn:
            conn.close()


def vertica_dedupe_test(
    table_name: str,
    unique_keys: Union[str, List[str]],
    conn: Optional[vertica_python.Connection] = None,
    conn_id: Optional[str] = None,
    date_col: Optional[str] = None,
    keep: str = "last",
    commit: bool = True,
    dir_path: Optional[str] = None,
    creds_path: Optional[str] = None,
):
    conn, close_conn = _resolve_vertica_test_conn(
        conn=conn,
        conn_id=conn_id,
        dir_path=dir_path,
        creds_path=creds_path,
    )
    try:
        return vertica_dedupe(
            table_name=table_name,
            unique_keys=unique_keys,
            conn=conn,
            date_col=date_col,
            keep=keep,
            commit=commit,
        )
    finally:
        if close_conn:
            conn.close()


def vertica_upsert_test(
    df: pd.DataFrame,
    table_name: str,
    unique_keys: Union[str, List[str]],
    conn=None,
    conn_id: Optional[str] = None,
    date_col: Optional[str] = None,
    days_back: Optional[int] = None,
    commit: bool = True,
    dir_path: Optional[str] = None,
    creds_path: Optional[str] = None,
):
    conn, close_conn = _resolve_vertica_test_conn(
        conn=conn,
        conn_id=conn_id,
        dir_path=dir_path,
        creds_path=creds_path,
    )
    try:
        return vertica_upsert(
            df=df,
            table_name=table_name,
            unique_keys=unique_keys,
            conn=conn,
            date_col=date_col,
            days_back=days_back,
            commit=commit,
        )
    finally:
        if close_conn:
            conn.close()
