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

import tdsbrondata

# --------------------------------------------------
# Helpers
# --------------------------------------------------

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)

# --------------------------------------------------
# Filtering & reduction
# --------------------------------------------------

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())
    )
    return dfCurrent.select(*[c["nameTable"] for c in columns])

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()

# --------------------------------------------------
# Source data
# --------------------------------------------------

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"

# --------------------------------------------------
# Mutation detection
# --------------------------------------------------

def extractMutations(sourceData, columns, dfExisting, filters):
    dfExistingReduced = reduceExistingData(dfExisting, columns)
    dfNewReduced = reduceNewData(applyFilters(sourceData, filters), columns)

    for c, t in dfExistingReduced.dtypes:
        dfNewReduced = dfNewReduced.withColumn(c, F.col(c).cast(t))

    return dfNewReduced.exceptAll(dfExistingReduced)

def extractDeletions(sourceData, columns, primaryKey, dfExisting, filters):
    dfExistingReduced = reduceExistingData(dfExisting, columns)
    dfNewReduced = reduceNewData(applyFilters(sourceData, filters), columns)

    if isinstance(primaryKey, list):
        cond = [dfExistingReduced[k] == dfNewReduced[k] for k in primaryKey]
    else:
        cond = dfExistingReduced[primaryKey] == dfNewReduced[primaryKey]

    return dfExistingReduced.join(dfNewReduced, cond, "left_anti")

# --------------------------------------------------
# SCD
# --------------------------------------------------

def processMerges(table, staging, columns, primaryKey, scdMoment):
    key_cond = (
        " AND ".join([f"t.{k} = s.{k}" for k in primaryKey])
        if isinstance(primaryKey, list)
        else f"t.{primaryKey} = s.{primaryKey}"
    )

    diff_cond = " OR ".join([
        f"ISNULL(t.{c['nameTable']},'') <> ISNULL(s.{c['nameTable']},'')"
        for c in columns
    ])

    cols = ", ".join(c["nameTable"] for c in columns)

    sql = f"""
    MERGE {table} t
    USING {staging} s
    ON {key_cond} AND t.CurrentFlag = 1

    WHEN MATCHED AND ({diff_cond})
    THEN UPDATE SET
        t.CurrentFlag = 0,
        t.ScdEndDate = {scdMoment}

    WHEN NOT MATCHED BY TARGET
    THEN INSERT ({cols}, CurrentFlag, ScdStartDate, ScdEndDate)
    VALUES ({cols}, 1, {scdMoment}, NULL);
    """

    executeSQL(sql)

def processDeletions(table, staging, primaryKey, scdMoment):
    if isinstance(primaryKey, list):
        key_cond = " AND ".join([f"t.{k} = s.{k}" for k in primaryKey])
        null_cond = " AND ".join([f"s.{k} IS NULL" for k in primaryKey])
    else:
        key_cond = f"t.{primaryKey} = s.{primaryKey}"
        null_cond = f"s.{primaryKey} IS NULL"

    sql = f"""
    UPDATE t
    SET
        t.CurrentFlag = 0,
        t.ScdEndDate = {scdMoment}
    FROM {table} t
    LEFT JOIN {staging} s
        ON {key_cond}
    WHERE
        {null_cond}
        AND t.CurrentFlag = 1;
    """
    executeSQL(sql)

# --------------------------------------------------
# Orchestration
# --------------------------------------------------

def processData(
    tableName,
    columns,
    primaryKey,
    filters,
    recreate,
    sourceDataOptions=[],
    sourceDataOverride=None,
    skipDeletions=False,
    scdMomentOverride=None
):
    sourceData = sourceDataOverride or getSourceData(tableName, sourceDataOptions)
    sourceData = sourceData.toDF(*[c.replace(" ", "_") for c in sourceData.columns])

    table = f"{tdsbrondata.schemaName}.{tableName}"
    staging = f"{tableName}_staging"

    ensureTable(table, columns)

    dfExisting = readFromTable(table)

    dfMutations = extractMutations(sourceData, columns, dfExisting, filters)
    dfDeletions = None if skipDeletions else extractDeletions(
        sourceData, columns, primaryKey, dfExisting, filters
    )

    writeToTable(dfMutations, staging, mode="overwrite")

    scdMoment = f"'{scdMomentOverride}'" if scdMomentOverride else "SYSUTCDATETIME()"

    processMerges(table, staging, columns, primaryKey, scdMoment)

    if dfDeletions is not None:
        processDeletions(table, staging, primaryKey, scdMoment)

    executeSQL(f"DROP TABLE IF EXISTS {staging}")

# --------------------------------------------------
# Full table write
# --------------------------------------------------

def writeFullTable(schemaName, tableName, df, columns, mode="overwrite"):

    if all(isinstance(c, str) for c in columns):
        columns = [{"nameTable": c, "nameSpark": c} for c in columns]
    
    table = f"{schemaName}.{tableName}"

    executeSQL(f"DROP TABLE IF EXISTS {table}")
    
    ensureTable(table, columns)
    
    rename = {c["nameSpark"]: c["nameTable"] for c in columns}
    for old, new in rename.items():
        if old != new and old in df.columns:
            df = df.withColumnRenamed(old, new)
    
    writeToTable(df, table, mode=mode)