# Copyright © 2024, UChicago Argonne, LLC
# BSD OPEN SOURCE LICENSE. Full license can be found in LICENSE.md
import re
from pathlib import Path

import pandas as pd
from polaris.utils.database.db_utils import list_tables_in_db
from polaris.utils.database.spatialite_utils import is_spatialite
from polaris.utils.dir_utils import mkdir_p
from polaris.utils.file_utils import df_to_file
from polaris.utils.signals import SIGNAL
from polaris.utils.structure_finder import find_table_fields

EXCL_NAME_PAT = [
    re.compile(e)
    for e in (
        "data_licenses",
        "geometry_columns",
        "spatial_ref_sys",
        "spatialite_history",
        "knn",
        "sql_statements_log",
        "sqlite_sequence",
        "views_geometry",
        "virts_geometry",
        "spatialindex",
        "elementarygeometries",
        "sqlite_stat",
        "table_metadata",
        "sequence",
        "metadata",
        "raster_coverages",
        "rl2map_configurations",
        "topologies",
        "^iso_metadata",
        "vector_coverages",
        "networks",
        "^idx_",
        "^se_",
        "^stored_",
        "^wms_",
    )
]


def dump_database_to_csv(
    conn, folder_name, signal=None, exclude_patterns=None, include_patterns=None, ext="csv"
) -> None:
    signal = signal or SIGNAL(object)
    exclude_patterns = exclude_patterns or EXCL_NAME_PAT

    folder_name = Path(folder_name)
    mkdir_p(folder_name)

    if is_spatialite(conn):
        geo_cols = get_table_data(conn, "geometry_columns")[
            ["f_table_name", "f_geometry_column", "geometry_type", "srid"]
        ]
        geo_cols.columns = ["table_name", "geo_column", "geometry_type", "srid"]
        geo_cols = geo_cols[geo_cols.table_name != "iso_metadata"]  # remove iso_metatable (standard Spatialite)
        geo_cols.to_csv(folder_name / "srids.csv", index=False)

    all_tables = [tn for tn in list_tables_in_db(conn) if not any(x.match(tn) for x in exclude_patterns)]
    if include_patterns is not None:
        include_patterns = [re.compile(p) if isinstance(p, str) else p for p in include_patterns]
        all_tables = [tn for tn in all_tables if any(x.match(tn) for x in include_patterns)]

    txt = "Dumping database"
    signal.emit(["start", "master", len(all_tables), txt])

    for table_name in all_tables:
        signal.emit(["start", "secondary", 2, table_name, txt])
        # get the schema
        schema = pd.read_sql_query(f"pragma table_info({table_name})", conn)
        schema.to_csv(folder_name / f"{table_name}.schema", index=False)

        # Get the table
        data = get_table_data(conn, table_name).round(decimals=6)
        signal.emit(["update", "secondary", 1, table_name, txt])
        df_to_file(data, folder_name / f"{table_name}.{ext}", index=False)
        signal.emit(["update", "secondary", 2, table_name, txt])

    signal.emit(["finished_dumping_procedure"])


def get_table_data(conn, table_name):
    fields, _, geo_field = find_table_fields(conn, table_name)
    fields = [f'"{x}"' for x in fields if '"' not in x]
    if geo_field is not None:
        fields.append(f'AsWKT("{geo_field}", 8) {geo_field}_wkt')
    keys = ",".join(fields)
    df = pd.read_sql_query(f"select {keys} from '{table_name}'", conn)
    df.columns = [x.lower() for x in df.columns]
    return df
