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

import tdsbrondata


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

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


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


def execute_sql(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 ensure_sql_scd_table(table, columns):
    """
    Creates the SCD2 table in SQL Server if it does not exist.
    Matches the structure expected by the SQL-based SCD logic.
    """

    column_defs = ",\n        ".join([
        f"{c['nameDelta']} NVARCHAR(255)" for c in columns
    ])

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


# --------------------------------------------------
# Filtering & reduction (UNCHANGED)
# --------------------------------------------------

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["nameDelta"] for c in columns])


def reduceNewData(df, columns):
    rename = {c["nameSpark"]: c["nameDelta"] 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 (UNCHANGED)
# --------------------------------------------------

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

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


# --------------------------------------------------
# SQL-BASED SCD2
# --------------------------------------------------

def merge_sql_scd2(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['nameDelta']},'') <> ISNULL(s.{c['nameDelta']},'')"
        for c in columns
    ])

    cols = ", ".join(c["nameDelta"] 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);
    """

    execute_sql(sql)


def delete_sql_scd2(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;
    """
    execute_sql(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"{table}_staging"

    # 🔹 ENSURE TABLE EXISTS (THE FIX)
    ensure_sql_scd_table(table, columns)

    dfExisting = read_sql_table(table)

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

    write_sql_table(dfMutations, staging, "overwrite")

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

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

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