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}
        );
    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):
    return df.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()

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]

    rowsToInsert = newData.join(
        existingData.select(*primaryKey),
        on=primaryKey,
        how="left_anti"
    )

    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} OR (n.{c} IS NULL AND e.{c} IS NOT NULL) OR (n.{c} IS NOT NULL AND e.{c} IS NULL))"
        for c in nonKeyColumns
    ])

    rowsToUpdate = joined.filter(changeCondition).select(
        *[F.col(f"n.{c}") for c in newData.columns]
    )

    return rowsToInsert, rowsToUpdate

def handleRowsToInsert(tableName, rowsToInsert, 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}) VALUES ({values})"
        if debug:
            print(sql)
        executeSQL(sql)

def handleRowsToUpdate(tableName, rowsToUpdate, primaryKey, debug):
    if isinstance(primaryKey, str):
        primaryKey = [primaryKey]

    if debug:
        print('-- Queries for update: ')

    if rowsToUpdate.count() == 0:
        return

    for row in rowsToUpdate.collect():
        rowDict = row.asDict()

        setClause = ", ".join([
            f"[{c}] = {escapeSqlValue(rowDict[c])}"
            for c in rowDict.keys() if c not in primaryKey
        ])

        whereClause = " AND ".join([
            f"[{k}] = {escapeSqlValue(rowDict[k])}"
            for k in primaryKey
        ])

        sql = f"UPDATE {tableName} SET {setClause} WHERE {whereClause}"

        if debug:
            print(sql)

        executeSQL(sql)

def processData(
    tableName,
    columns,
    primaryKey,
    filters,
    sourceDataOptions=[],
    sourceDataOverride=None,
    debug=False
):
    table = f"{tdsbrondata.schemaName}.{tableName}"

    ensureTable(table, columns)

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

    existingData = readFromTable(table)
    existingDataReduced = reduceExistingData(existingData, columns)

    rowsToInsert, rowsToUpdate = extractMutations(
        newDataReduced,
        existingDataReduced,
        primaryKey
    )

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