import sys
import os
from functools import reduce
from pyspark.sql import functions as F
from pyspark.sql.types import *

import tdsbrondata

def readFromTable(table):
    return (
        tdsbrondata._spark.read
            .format("jdbc")
            .option("url", tdsbrondata.jdbcUrl)
            .option("dbtable", table)
            .option("accessToken", tdsbrondata.token)
            .load()
    )

def writeToTable(df, table, mode="overwrite"):
    (
        df.write
          .format("jdbc")
          .option("url", tdsbrondata.jdbcUrl)
          .option("dbtable", table)
          .option("accessToken", tdsbrondata.token)
          .mode(mode)
          .save()
    )

def executeSQL(sql):
    jvm = tdsbrondata._spark._sc._jvm

    props = jvm.java.util.Properties()
    props.setProperty("accessToken", tdsbrondata.token)

    conn = jvm.java.sql.DriverManager.getConnection(
        tdsbrondata.jdbcUrl,
        props
    )

    stmt = conn.createStatement()
    stmt.execute(sql)
    stmt.close()
    conn.close()

def ensureTable(table, columns):
    columnDefs = ",\n        ".join([
        f"[{c['nameTable']}] NVARCHAR(255)" for c in columns
    ])

    sql = f"""
    IF OBJECT_ID('{table}', 'U') IS NULL
    BEGIN
        CREATE TABLE {table} (
            {columnDefs},
            SurrogateKey BIGINT IDENTITY(1,1) NOT NULL,
            CurrentFlag BIT NOT NULL,
            ScdStartDate DATETIME2 NOT NULL,
            ScdEndDate DATETIME2 NULL
        );
    END
    """
    executeSQL(sql)

def escapeSqlValue(v):
    if v is None:
        return "NULL"
    if isinstance(v, str):
        v_escaped = v.replace("'", "''")
        return f"'{v_escaped}'"
    return str(v)

def applyFilters(df, filters):
    if not filters:
        return df

    conditions = reduce(
        lambda acc, f: acc & (F.col(f[0]) == f[2]) if f[1] == "==" else acc & (F.col(f[0]) != f[2]),
        filters,
        F.lit(True)
    )
    return df.filter(conditions)

def reduceExistingData(df, columns):
    dfCurrent = df.filter(
        (F.col("CurrentFlag") == 1) & (F.col("ScdEndDate").isNull())
    )
    selectCols = [c["nameTable"] for c in columns] + ["SurrogateKey"]
    return dfCurrent.select(*selectCols)

def reduceNewData(df, columns):
    rename = {c["nameSpark"]: c["nameTable"] for c in columns}
    df = df.select(*rename.keys())
    for old, new in rename.items():
        df = df.withColumnRenamed(old, new)
    return df.distinct()

def getSourceData(tableName, sourceDataOptions):
    if tdsbrondata.sourceDataMode == "automatic":
        return getAutomaticSourceData(tableName)
    return getManualSourceData(tableName, sourceDataOptions)

def getAutomaticSourceData(tableName):
    full = getAutomaticSourceDataPath(tableName, "full")
    incr = getAutomaticSourceDataPath(tableName, "incremental")

    if tdsbrondata._notebookutils.fs.exists(full):
        return tdsbrondata._spark.read.parquet(full)
    if tdsbrondata._notebookutils.fs.exists(incr):
        return tdsbrondata._spark.read.parquet(incr)

    raise FileNotFoundError(tableName)

def getManualSourceData(tableName, sourceDataOptions):
    path = getManualSourceDataPath(tableName)
    reader = tdsbrondata._spark.read
    if sourceDataOptions:
        for opt, val in sourceDataOptions:
            reader = reader.option(opt, val)
    return reader.csv(path)

def getAutomaticSourceDataPath(tableName, loadType):
    path = f"{tdsbrondata.automaticDataPath}/{tableName}_{loadType}.parquet"
    path = path.replace(tdsbrondata.lakehouseNameSilver, tdsbrondata.lakehouseNameBronze)
    path = path.replace(tdsbrondata.schemaNameSilver, tdsbrondata.schemaNameBronze)
    return path

def getManualSourceDataPath(tableName):
    return f"{tdsbrondata.manualDataPath}/{tdsbrondata.sourceDataPeriod}_{tableName}.csv"

def extractMutations(newData, existingData, primaryKey):
    if isinstance(primaryKey, str):
        primaryKey = [primaryKey]

    # Inserts: data that is in newData but missing in existingData
    rowsToInsert = newData.join(
        existingData.select(*primaryKey),
        on=primaryKey,
        how="left_anti"
    )

    # Updates: data that is in newData and in existingData but with changes
    joined = newData.alias("n").join(
        existingData.alias("e"),
        on=primaryKey,
        how="inner"
    )
    nonKeyColumns = [c for c in newData.columns if c not in primaryKey]
    changeCondition = " OR ".join([f"n.{c} != e.{c}" for c in nonKeyColumns])
    rowsToUpdate = joined.filter(changeCondition).select(
        *primaryKey,
        *[F.col(f"n.{c}") for c in nonKeyColumns],
        F.col("e.SurrogateKey").alias("SurrogateKey"),
    )

    # Deletes: data that is in existingData but missing in newData
    rowsToDelete = existingData.join(
        newData.select(*primaryKey),
        on=primaryKey,
        how="left_anti"
    )

    return rowsToInsert, rowsToUpdate, rowsToDelete

def handleRowsToInsert(tableName, rowsToInsert, scdMoment, debug):
    if debug:
        print('-- Queries for insert: ')

    if rowsToInsert.count() == 0:
        return

    for row in rowsToInsert.collect():
        columns = ", ".join([f"[{c}]" for c in row.asDict().keys()])
        values = ", ".join([escapeSqlValue(v) for v in row])
        sql = f"INSERT INTO {tableName} ({columns}, [CurrentFlag], [ScdStartDate], [ScdEndDate]) " \
              f"VALUES ({values}, 1, {scdMoment}, NULL)"
        if debug:
            print(sql)
        executeSQL(sql)

def handleRowsToUpdate(tableName, rowsToUpdate, scdMoment, debug):
    if debug:
        print('-- Queries for update: ')

    if rowsToUpdate.count() == 0:
        return

    surrogateKeys = [str(row.SurrogateKey) for row in rowsToUpdate.select("SurrogateKey").collect()]
    if surrogateKeys:
        keys_str = ", ".join(surrogateKeys)
        sql_expire = f"UPDATE {tableName} SET [CurrentFlag] = 0, [ScdEndDate] = {scdMoment} " \
                     f"WHERE [SurrogateKey] IN ({keys_str})"
        if debug:
            print(sql_expire)
        executeSQL(sql_expire)

    for row in rowsToUpdate.collect():
        columns = [c for c in row.asDict().keys() if c != "SurrogateKey"]
        columns_str = ", ".join([f"[{c}]" for c in columns])
        values_str = ", ".join([escapeSqlValue(row[c]) for c in columns])
        sql_insert = f"INSERT INTO {tableName} ({columns_str}, [CurrentFlag], [ScdStartDate], [ScdEndDate]) " \
                     f"VALUES ({values_str}, 1, {scdMoment}, NULL)"
        if debug:
            print(sql_insert)
        executeSQL(sql_insert)

def handleRowsToDelete(tableName, rowsToDelete, scdMoment, debug):
    if debug:
        print('-- Queries for delete: ')

    if rowsToDelete.count() == 0:
        return

    surrogateKeys = [str(row.SurrogateKey) for row in rowsToDelete.select("SurrogateKey").collect()]
    if surrogateKeys:
        keys_str = ", ".join(surrogateKeys)
        sql = f"UPDATE {tableName} SET [CurrentFlag] = 0, [ScdEndDate] = {scdMoment} " \
              f"WHERE [SurrogateKey] IN ({keys_str})"
        if debug:
            print(sql)
        executeSQL(sql)

def processData(
    tableName,
    columns,
    primaryKey,
    filters,
    sourceDataOptions=[],
    sourceDataOverride=None,
    skipDeletions=False,
    scdMomentOverride=None,
    debug=False
):

    # Determine the table name with the correct schema
    table = f"{tdsbrondata.schemaName}.{tableName}"

    # Create the table if it didn't exist yet
    ensureTable(table, columns)

    # Determine the SCD moment
    scdMoment = f"'{scdMomentOverride}'" if scdMomentOverride else "SYSUTCDATETIME()"

    # New data
    newData = sourceDataOverride or getSourceData(tableName, sourceDataOptions)
    newData = newData.toDF(*[c.replace(" ", "_") for c in newData.columns])
    newDataReduced = reduceNewData(applyFilters(newData, filters), columns)

    # Existing data
    existingData = readFromTable(table)
    existingDataReduced = reduceExistingData(existingData, columns)

    # Get mutations
    rowsToInsert, rowsToUpdate, rowsToDelete = extractMutations(
        newDataReduced,
        existingDataReduced,
        primaryKey
    )

    if not skipDeletions:
        handleRowsToDelete(table, rowsToDelete, scdMoment, debug)

    handleRowsToUpdate(table, rowsToUpdate, scdMoment, debug)
    handleRowsToInsert(table, rowsToInsert, scdMoment, debug)