import tempfile
import pandas as pd
from pathlib import Path
from psycopg2 import connect
import psycopg2
from xplordb.datamodel.person import Person
from xplordb.datamodel.company import Company
from xplordb.datamodel.dataset import Dataset
from xplordb.import_data import ImportData
from xplordb.import_ddl import ImportDDL
from xplordb.sqlalchemy.base import create_session


path = Path(tempfile.mkdtemp())


collar_columns = ["data_set", "hole_id", "x", "y", "z", "eoh", "planned_x", "planned_y", "planned_z", "planned_eoh", "dip", "azimuth", "loaded_by", "load_date", "srid", "project_srid"]
survey_columns = ["data_set", "hole_id", "depth_m", "dip", "azimuth_grid", "loaded_by"]
categories_columns = ["table_name", "name", "name_column", "schema"]
person_columns = ["code", "person", "load_date"]
dataset_columns = ["data_set", "full_name", "loaded_by", "load_date"]
profile_columns = ["id", "name", "type", "assay_name", "parameters"]
auto_columns = ["id", "name", "profile", "source_folder", "person", "dataset", "frequency", "log_folder", "creation_date", "last_run_date", "run_count", "failed_run_count"]
assay_column_columns = ["name", "assay", "series_type", "unit", "category_name", "upper_box_column", "lower_box_column", "upper_whisker_column", "lower_whisker_column", 
               "image_format_col", "detection_min_col", "detection_max_col", "display_name", "symbology"]
assay_column_definition_columns = ["name", "variable", "col"]
assay_definition_columns = ["variable", "table_name", "hole_id_col", "x_col", "dataset_col", "person_col", "import_date_col", "x_end_col", "schema"]
assays_columns = ["variable", "domain", "data_extent", "display_name"]


def migrate(conn) -> None:

    # 1st step : export data as csv
    # collars
    with conn.cursor() as cur:
        cur.execute(f"select {','.join(collar_columns)} from dh.collar")
        res = cur.fetchall()
        df = pd.DataFrame(res, columns= collar_columns)
        df.to_csv(path / "collars.csv", index=False)

    # metadata
    with conn.cursor() as cur:
        cur.execute("""
                    SELECT column_name
                    FROM information_schema.columns
                    WHERE table_schema = 'dh'
                    AND table_name = 'metadata';
                    """)
        columns = cur.fetchall()
        columns = [col[0] for col in columns]

        cur.execute(f"select {','.join(columns)} from dh.metadata")
        res = cur.fetchall()
        df = pd.DataFrame(res, columns= columns)
        df.to_csv(path / "metadata.csv", index=False)

    # metadata types
    with conn.cursor() as cur:
        cur.execute("""
                    SELECT column_name, data_type
                    FROM information_schema.columns
                    WHERE table_schema = 'dh'
                    AND table_name = 'metadata';
                    """)
        types = cur.fetchall()
        df = pd.DataFrame(types, columns= ["column", "type"])
        df.to_csv(path / "metadata_type.csv", index=False)

    # surveys
    with conn.cursor() as cur:
        cur.execute(f"select {','.join(survey_columns)} from dh.surv")
        res = cur.fetchall()
        df = pd.DataFrame(res, columns= survey_columns)
        df.to_csv(path / "surveys.csv", index=False)

    # categories
    with conn.cursor() as cur:
        
        cur.execute(f"select {','.join(categories_columns)} from ref.categories")
        res = cur.fetchall()
        df = pd.DataFrame(res, columns= categories_columns)
        df.to_csv(path / "categories.csv", index=False)

    # person
    with conn.cursor() as cur:
        
        cur.execute(f"select {','.join(person_columns)} from ref.person")
        res = cur.fetchall()
        df = pd.DataFrame(res, columns= person_columns)
        df.to_csv(path / "person.csv", index=False)

    # dataset
    with conn.cursor() as cur:
        
        cur.execute(f"select {','.join(dataset_columns)} from ref.data_sets")
        res = cur.fetchall()
        df = pd.DataFrame(res, columns= dataset_columns)
        df.to_csv(path / "datasets.csv", index=False)

    # import profile
    with conn.cursor() as cur:
        
        cur.execute(f"select {','.join(profile_columns)} from ref.import_profiles")
        res = cur.fetchall()
        df = pd.DataFrame(res, columns= profile_columns)
        df.to_csv(path / "import_profiles.csv", index=False)

    # automated imports
    with conn.cursor() as cur:
        
        cur.execute(f"select {','.join(auto_columns)} from ref.automated_imports")
        res = cur.fetchall()
        df = pd.DataFrame(res, columns= auto_columns)
        df.to_csv(path / "automated_imports.csv", index=False)

    # assay metadatas
    with conn.cursor() as cur:
        
        cur.execute(f"select {','.join(assay_column_columns)} from assay.assay_column")
        res = cur.fetchall()
        df = pd.DataFrame(res, columns= assay_column_columns)
        df.to_csv(path / "assay_column.csv", index=False)

    with conn.cursor() as cur:
        
        cur.execute(f"select {','.join(assay_column_definition_columns)} from assay.assay_column_definition")
        res = cur.fetchall()
        df = pd.DataFrame(res, columns= assay_column_definition_columns)
        df.to_csv(path / "assay_column_definition.csv", index=False)

    with conn.cursor() as cur:
        
        cur.execute(f"select {','.join(assay_definition_columns)} from assay.assay_definition")
        res = cur.fetchall()
        df = pd.DataFrame(res, columns= assay_definition_columns)
        df.to_csv(path / "assay_definition.csv", index=False)

    with conn.cursor() as cur:
        
        cur.execute(f"select {','.join(assays_columns)} from assay.assays")
        res = cur.fetchall()
        df = pd.DataFrame(res, columns= assays_columns)
        df.to_csv(path / "assays.csv", index=False)


    # assay raw tables
    assays = pd.read_csv(path / "assay_definition.csv")
    for assay in assays["table_name"]:
        with conn.cursor() as cur:
            cur.execute(f"""
                        SELECT column_name
                        FROM information_schema.columns
                        WHERE table_schema = 'assay'
                        AND table_name = '{assay}';
                        """)
            columns = cur.fetchall()
            columns = [col[0] for col in columns]
            # put column name in double quotes
            quoted_columns = [f'"{column}"' for column in columns]

            cur.execute(f"select {','.join(quoted_columns)} from assay.{assay}")
            res = cur.fetchall()
            df = pd.DataFrame(res, columns= columns)
            filename = "assay_table_" + assay + ".csv"
            df.to_csv(path / filename, index=False)

    # category tables
    categories = pd.read_csv(path / "categories.csv")
    for category in categories["table_name"]:
        with conn.cursor() as cur:
            columns = ["name", "description"]
            cur.execute(f"select {','.join(columns)} from public.{category}")
            res = cur.fetchall()
            df = pd.DataFrame(res, columns= columns)
            filename = "category_table_" + category + ".csv"
            df.to_csv(path / filename, index=False)


    # 2nd step : delete current schema (except public)
    q = """
        SELECT schema_name
        FROM information_schema.schemata
        WHERE schema_name NOT LIKE 'pg_%'
        AND schema_name != 'information_schema';
        """

    with conn.cursor() as cur:
        cur.execute(q)
        schemas = cur.fetchall()
        schemas = [schema[0] for schema in schemas]

    for schema in schemas:
        if schema == "public":
            continue
        with conn.cursor() as cur:
            cur.execute(f"DROP SCHEMA {schema} CASCADE;")
        conn.commit()


    # 3rd step : install new schema
    new_database = "xplordb_migrated"

    import_ddl = ImportDDL(db=new_database,
                        user='postgres',
                        password='larmiche',
                        host='localhost',
                        port=5432)
    import_ddl.import_xplordb_schema()

    # 4th step : import data
    # ref : company, persons, datasets, import_profile, automated_import, categories

    session = create_session(
        database=new_database,
        host='localhost',
        port=5432,
        user='postgres',
        password='larmiche')

    company = Company("openlog")
    import_data = ImportData(session)
    import_data.import_companies_array([company])

    person_df = pd.read_csv(path / "person.csv")
    dataset_df = pd.read_csv(path / "datasets.csv")

    persons = [Person(name=person_name, company=company) for person_name in person_df["person"]]
    trigrams = {row["code"]:row["person"] for _, row in person_df.iterrows()}

    datasets = []
    for _, row in dataset_df.iterrows():
        person_trigram = row["loaded_by"]
        ds_name = row["full_name"]
        person = Person(name=trigrams.get(person_trigram), company=company)
        dataset = Dataset(name=ds_name, person=person)
        datasets.append(dataset)

    import_data.import_persons_array(persons)
    import_data.import_datasets_array(datasets)
    session.commit()

    # new person, dataset IDs
    person_df = pd.read_csv(path / "person.csv")
    person_dict = {}
    for i, person in person_df.iterrows():
        person_dict[person["code"]] = i + 1

    dataset_df = pd.read_csv(path / "datasets.csv")
    dataset_dict = {}
    for i, dataset in dataset_df.iterrows():
        dataset_dict[dataset["data_set"]] = i + 1

    # import profiles
    profiles = pd.read_csv(path / "import_profiles.csv")
    with conn.cursor() as cur:
        for _, profile in profiles.iterrows():
            quoted_columns = [f'"{column}"' for column in profile_columns]
            col_str = ",".join(quoted_columns)
            cur.execute(f"INSERT INTO ref.import_profiles({col_str}) VALUES({','.join(['%s' for col in profile_columns])})",
                        profile.astype(str).to_list())

    conn.commit()

    # automated imports
    autos = pd.read_csv(path / "automated_imports.csv")
    with conn.cursor() as cur:
        for _, auto in autos.iterrows():
            # change person and dataset by IDs
            auto["person"] = person_dict.get(auto["person"])
            auto["dataset"] = dataset_dict.get(auto["dataset"])
            row = auto.to_list()
            row = [str(elt) if not pd.isna(elt) else None for elt in row]

            quoted_columns = [f'"{column}"' for column in auto_columns]
            col_str = ",".join(quoted_columns)
            cur.execute(f"INSERT INTO ref.automated_imports({col_str}) VALUES({','.join(['%s' for col in auto_columns])})",
                        row)
    conn.commit()

    # collars
    collar_df = pd.read_csv(path / "collars.csv")
    with conn.cursor() as cur:
        for _, collar in collar_df.iterrows():
            # change person and dataset by IDs
            collar["loaded_by"] = person_dict.get(collar["loaded_by"])
            collar["data_set"] = dataset_dict.get(collar["data_set"])
            quoted_columns = [f'"{column}"' for column in collar_columns]
            col_str = ",".join(quoted_columns)
            cur.execute(f"INSERT INTO dh.collar({col_str}) VALUES({','.join(['%s' for col in collar_columns])})",
                        collar.astype(str).to_list())
    conn.commit()

    # surveys
    survey_df = pd.read_csv(path / "surveys.csv")
    with conn.cursor() as cur:
        for _, survey in survey_df.iterrows():
            # change person and dataset by IDs
            survey["loaded_by"] = person_dict.get(survey["loaded_by"])
            survey["data_set"] = dataset_dict.get(survey["data_set"])
            quoted_columns = [f'"{column}"' for column in survey_columns]
            col_str = ",".join(quoted_columns)
            cur.execute(f"INSERT INTO dh.surv({col_str}) VALUES({','.join(['%s' for col in survey_columns])})",
                        survey.astype(str).to_list())
    conn.commit()

    # metadata (add columns)
    metadata_df = pd.read_csv(path / "metadata.csv")
    metadata_types = pd.read_csv(path / "metadata_type.csv")
    new_columns = [column for column in metadata_df.columns if column not in  ["hole_id"]]
    with conn.cursor() as cur:
        for new_column in new_columns:
            type_ = metadata_types.loc[metadata_types["column"] == new_column,"type"].values[0]
            cur.execute(f"ALTER TABLE dh.metadata ADD COLUMN {new_column} {type_}")

    conn.commit()

    # metadata (update)
    with conn.cursor() as cur:
        for new_column in new_columns:
            for _, metadata in metadata_df.iterrows():
                hole_id = metadata["hole_id"]
                value = metadata[new_column]
                if pd.isna(value):
                    continue
                cur.execute(f'UPDATE dh.metadata SET "{new_column}" = %s WHERE hole_id = %s',
                            (value, hole_id))

    conn.commit()

    # categories (metadata)
    categories_df = pd.read_csv(path / "categories.csv")
    categories_df = categories_df.replace(pd.NA, None)
    with conn.cursor() as cur:
        for _, category in categories_df.iterrows():
            row = category.to_list()
            row = [elt if not pd.isna(elt) else "" for elt in row]
            cur.execute("INSERT INTO ref.categories VALUES (%s, %s, %s, %s)",
                        row)

    conn.commit()

    # assay (metadatas)
    assays_df = pd.read_csv(path / "assays.csv")
    with conn.cursor() as cur:
        for _, assay in assays_df.iterrows():
            row = assay.to_list()
            cur.execute("INSERT INTO assay.assays VALUES (%s, %s, %s, %s)",
                        row)
    conn.commit()

    assay_def_df = pd.read_csv(path / "assay_definition.csv")
    with conn.cursor() as cur:
        for _, assay in assay_def_df.iterrows():
            row = assay.to_list()
            cur.execute("INSERT INTO assay.assay_definition VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)",
                        row)
    conn.commit()

    assay_col_df = pd.read_csv(path / "assay_column.csv")
    with conn.cursor() as cur:
        for _, assay in assay_col_df.iterrows():
            row = assay.to_list()
            row = [elt if not pd.isna(elt) else "" for elt in row]
            cur.execute("INSERT INTO assay.assay_column VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)",
                        row)
    conn.commit()

    assay_col_def_df = pd.read_csv(path / "assay_column_definition.csv")
    with conn.cursor() as cur:
        for _, assay in assay_col_def_df.iterrows():
            row = assay.to_list()
            cur.execute("INSERT INTO assay.assay_column_definition VALUES (%s, %s, %s)",
                        row)
    conn.commit()

    # category tables (no need to create them because public schema is preserved)

    # assay tables creation

    def sql_assay_table_definition(table_name: str):
        assay = assays_df.loc[assays_df["variable"] == table_name].iloc[0]
        domain = assay["domain"]
        extent = assay["data_extent"]

        # columns description
        columns = assay_col_df.loc[assay_col_df["assay"] == table_name]

        # raw table
        df = pd.read_csv(path / str("assay_table_" + table_name + ".csv"))
        col_sql = []
        for col in df.columns:
            
            if col == "hole":
                type_ = "character varying NOT NULL"
            elif col == "dataset":
                type_ = "integer NOT NULL"
            elif col == "person":
                type_ = "integer REFERENCES ref.person(id) ON DELETE CASCADE"
            elif col == "import_date":
                type_ = "DATE"
            elif col == "x":
                if domain == "depth":
                    type_ = "FLOAT NOT NULL"
                else:
                    type_ = "timestamp without time zone NOT NULL"

            elif col == "x_end":
                if domain == "depth":
                    type_ = "FLOAT NOT NULL"
                else:
                    type_ = "timestamp without time zone NOT NULL"
            
            # measure column
            elif col in columns["name"].to_list():
                column = columns.loc[columns["name"] == col].iloc[0]
                serie_type = column["series_type"]

                if serie_type == "numerical":
                    type_ = "REAL"
                elif serie_type == "polar":
                    type_ = "assay.azimuth"
                elif serie_type in ["nominal", "categorical"]:
                    type_ = "TEXT"
                elif serie_type == "datetime":
                    type_ = "DATE"
                elif serie_type == "spherical":
                    type_ = "assay.spherical_data"
                elif serie_type == "imagery":
                    type_ = "BLOB"
            
            # detection limits or uncertainties
            else:
                type_ = "REAL"
            
            col_sql.append(f'"{col}"' + " " + type_)

        # primary key
        pkey_sql = "PRIMARY KEY (hole, dataset, x)" if extent == "discrete" else "PRIMARY KEY (hole, dataset, x, x_end)"

        # foreign key
        key_sql = "FOREIGN KEY (hole, dataset) REFERENCES dh.collar(hole_id, data_set) ON DELETE CASCADE"

        col_sql.append(pkey_sql)
        col_sql.append(key_sql)
        sql = f'CREATE TABLE assay."{table_name}" ({",".join(col_sql)});'
        return sql

    # get (hole_id, data_set) pairs from dh.collar
    with conn.cursor() as cur:
        cur.execute("select hole_id, data_set from dh.collar")
        collar_pkeys = cur.fetchall()

    # insert data
    with conn.cursor() as cur:
        for _, assay in assays_df.iterrows():
        
            table_name = assay["variable"]
            # create table
            cur.execute(sql_assay_table_definition(table_name))

            # grant permissions to xdb roles
            cur.execute(f'ALTER TABLE assay."{table_name}" OWNER TO xdb_admin;')
            cur.execute(f'GRANT SELECT, TRIGGER ON TABLE assay."{table_name}" TO xdb_viewer;')
            cur.execute(f'GRANT SELECT, UPDATE, INSERT, TRIGGER ON TABLE assay."{table_name}" TO xdb_logger;')
            cur.execute(f'GRANT ALL ON TABLE assay."{table_name}" TO xdb_admin;')


            print("Importing assay table : ", table_name)
            # insert
            df = pd.read_csv(path / str("assay_table_" + table_name + ".csv"))
            placeholders = ",".join(["%s" for i in range(df.shape[1])])
            for _, row in df.iterrows():
                # replace dataset and person by IDs
                row["person"] = person_dict.get(row["person"])
                row["dataset"] = dataset_dict.get(row["dataset"])
                hole_dataset = (row["hole"], row["dataset"])
                row = [elt if not pd.isna(elt) else None for elt in row]
                # some assay have hole_id not in dh.collar !
                if hole_dataset in collar_pkeys:
                    cur.execute(f'INSERT INTO assay."{table_name}" VALUES ({placeholders})',
                            row)
    conn.commit()