"""許可済みドメインを、ドメイン単位でマッピングして検出結果から除外する。"""

from pyspark.sql import Row, functions as F, types as T
import csv
import json
import re

WHITELIST_TABLE = "domain_whitelist"
SOURCE_RESULT_TABLE = "input_allp_audit_trn_spark_driver_log_line_detected_result"
TARGET_TABLE = "input_allp_audit_trn_spark_driver_log_line_detected_result_after_exclusion"
MAPPING_AUDIT_TABLE = "input_allp_audit_trn_spark_driver_log_line_exclusion_mapping"
RESULT_AUDIT_TABLE = "input_allp_audit_trn_spark_driver_log_line_exclusion_result_audit"

# 動作確認時はサンプル DataFrame を使う。本番の管理テーブルを使う場合は False に変更する。
USE_SAMPLE_WHITELIST = True

sample_whitelist_rows = [
    Row(
        notebook_id="03e46dd6-eb7a-493d-bfec-43a49de217e6",
        domain_name='"*.spark*triprodje.dfs.core.windows.net","*.pbidedicated.windows.net","exec.japaneast.notebook.windows.net","olst*.dfs.core.windows.net","operation-service","tokenservice*.japaneast.trident.azuresynapse.net"',
        workspace_id="4a37edb7-95af-48ee-bf31-b0719c9f5efc",
        exclusion_date="2026-09-11",
    ),
]

if USE_SAMPLE_WHITELIST:
    domain_whitelist_df = spark.createDataFrame(sample_whitelist_rows)
else:
    domain_whitelist_df = spark.table(WHITELIST_TABLE)

# 管理テーブルは 1 行の domain_name に複数のドメインを保持する。
whitelist_distinct = (
    domain_whitelist_df
    .select(
        F.col("notebook_id").alias("アイテム物理名"),
        F.col("domain_name").alias("ドメイン名"),
        F.col("workspace_id").alias("所属ワークスペース"),
        F.to_date("exclusion_date", "yyyy-MM-dd").alias("除外適用日"),
    )
    .distinct()
)
display(whitelist_distinct.orderBy("所属ワークスペース", "アイテム物理名", "ドメイン名"))

def normalize_domain_token(value):
    """空白、引用符、JSON 配列の角括弧だけを除去してドメインパターンを残す。"""
    if value is None:
        return None
    token = str(value).strip().strip("[] \t\r\n\"'").strip()
    return token or None

def split_domain_string(value):
    """利用者が入力し得る domain_name の 4 形式を配列に変換する。"""
    if value is None:
        return []
    if isinstance(value, (list, tuple)):
        return [token for token in (normalize_domain_token(item) for item in value) if token]

    text = str(value).strip()
    if not text:
        return []

    # 形式 1: ["domain-a", "domain-b"] の JSON 配列
    try:
        parsed = json.loads(text)
        if isinstance(parsed, list):
            return [token for token in (normalize_domain_token(item) for item in parsed) if token]
    except (TypeError, ValueError):
        pass

    # 形式 2: "domain-a","domain-b"、形式 3: domain-a, domain-b、
    # 形式 4: 1 row に 1 domain のいずれもここで処理する。
    try:
        fields = next(csv.reader([text], skipinitialspace=True))
    except (csv.Error, StopIteration):
        fields = text.split(",")
    return [token for token in (normalize_domain_token(item) for item in fields) if token]

def parse_target_string(value):
    """[\"domain-a\",\"domain-b\"] 形式の JSON 文字列を配列に変換する。"""
    if not value:
        return []
    try:
        parsed = json.loads(value)
        return [token for token in (normalize_domain_token(item) for item in parsed) if token] if isinstance(parsed, list) else []
    except (TypeError, ValueError):
        return []

def wildcard_matches(target, allowed_pattern):
    """許可パターンの * を任意文字列として、対象ドメインと完全一致で照合する。"""
    if not target or not allowed_pattern:
        return False
    regex = "^" + re.escape(allowed_pattern.strip().lower()).replace(r"\*", ".*") + "$"
    return re.fullmatch(regex, target.strip().lower()) is not None

string_array_type = T.ArrayType(T.StringType(), False)
split_domain_udf = F.udf(split_domain_string, string_array_type)
parse_targets_udf = F.udf(parse_target_string, string_array_type)
wildcard_matches_udf = F.udf(wildcard_matches, T.BooleanType())

# whitelist は照合のためだけに一時展開する。管理テーブル自体の 1 行構造は変更しない。
whitelist_expanded = (
    domain_whitelist_df
    .select(
        "notebook_id", "workspace_id",
        F.to_date("exclusion_date", "yyyy-MM-dd").alias("exclusion_date"),
        F.explode_outer(split_domain_udf("domain_name")).alias("whitelist_domain_pattern"),
    )
    .where(F.col("whitelist_domain_pattern").isNotNull())
    .distinct()
)

source_result = spark.table(SOURCE_RESULT_TABLE)
GROUP_KEYS = ["workspace_id", "notebook_id", "livy_id", "submitted_datetime"]

# non_fabric_targets が旧配列型の場合にも動作するよう、いったん JSON 文字列に統一する。
target_field = next(field for field in source_result.schema.fields if field.name == "non_fabric_targets")
if isinstance(target_field.dataType, T.ArrayType):
    source_with_targets = source_result.withColumn("_non_fabric_targets_json", F.to_json("non_fabric_targets"))
else:
    source_with_targets = source_result.withColumn("_non_fabric_targets_json", F.col("non_fabric_targets").cast("string"))

# 検出結果の JSON 文字列もドメイン単位に一時展開する。
target_expanded = (
    source_with_targets
    .select(
        *GROUP_KEYS,
        F.to_date("submitted_datetime").alias("submitted_date"),
        F.explode_outer(parse_targets_udf("_non_fabric_targets_json")).alias("detected_domain"),
    )
    .where(F.col("detected_domain").isNotNull())
    .distinct()
)

# workspace、notebook、適用日が一致する whitelist だけを候補として結合し、
# その後 wildcard を使ってドメイン単位で照合する。
mapping_candidates = target_expanded.join(
    F.broadcast(whitelist_expanded),
    on=(
        (target_expanded.workspace_id == whitelist_expanded.workspace_id)
        & (target_expanded.notebook_id == whitelist_expanded.notebook_id)
        & (target_expanded.submitted_date >= whitelist_expanded.exclusion_date)
    ),
    how="left",
)

domain_mapping = (
    mapping_candidates
    .withColumn("_pattern_matches", wildcard_matches_udf("detected_domain", "whitelist_domain_pattern"))
    .groupBy(*[target_expanded[key] for key in GROUP_KEYS], "detected_domain")
    .agg(
        F.max(F.when(F.col("_pattern_matches"), F.lit(1)).otherwise(F.lit(0))).alias("_is_excluded"),
        F.sort_array(F.collect_set(F.when(F.col("_pattern_matches"), F.col("whitelist_domain_pattern")))).alias("一致した許可ドメイン"),
        F.sort_array(F.collect_set(F.when(F.col("_pattern_matches"), F.col("exclusion_date").cast("string")))).alias("適用した除外日"),
    )
    .withColumn("除外可否", F.when(F.col("_is_excluded") == 1, F.lit("除外")).otherwise(F.lit("対象外")))
    .drop("_is_excluded")
)

# 1 件でも許可されていないドメインが残る場合は、その row を最終結果に残す。
mapping_group = (
    domain_mapping
    .groupBy(*GROUP_KEYS)
    .agg(
        F.sort_array(F.collect_set("detected_domain")).alias("_all_targets"),
        F.sort_array(F.collect_set(F.when(F.col("除外可否") == "除外", F.col("detected_domain")))).alias("_excluded_targets"),
        F.sort_array(F.collect_set(F.when(F.col("除外可否") == "対象外", F.col("detected_domain")))).alias("_remaining_targets"),
        F.sort_array(F.array_distinct(F.flatten(F.collect_list("一致した許可ドメイン")))).alias("_matched_whitelist_domains"),
    )
)

# ドメイン単位のマッピング結果を監査用テーブルに保存する。
(domain_mapping.write.format("delta").mode("overwrite").option("overwriteSchema", "true").saveAsTable(MAPPING_AUDIT_TABLE))

empty_string_array = F.array().cast("array<string>")
exclusion_evaluated = (
    source_with_targets
    .join(mapping_group, GROUP_KEYS, "left")
    .withColumn("_all_targets", F.coalesce(F.col("_all_targets"), empty_string_array))
    .withColumn("_excluded_targets", F.coalesce(F.col("_excluded_targets"), empty_string_array))
    .withColumn("_remaining_targets", F.coalesce(F.col("_remaining_targets"), empty_string_array))
    .withColumn(
        "除外結果",
        F.when(F.size("_all_targets") == 0, F.lit("除外対象なし"))
         .when(F.size("_remaining_targets") == 0, F.lit("全対象を除外"))
         .when(F.size("_excluded_targets") == 0, F.lit("除外対象なし"))
         .otherwise(F.lit("一部対象を除外")),
    )
    .withColumn("non_fabric_targets", F.to_json("_remaining_targets"))
    .withColumn("除外したドメイン", F.to_json("_excluded_targets"))
    .withColumn("一致した許可ドメイン", F.to_json(F.coalesce(F.col("_matched_whitelist_domains"), empty_string_array)))
    .drop("_non_fabric_targets_json", "_all_targets", "_excluded_targets", "_remaining_targets", "_matched_whitelist_domains")
)

(exclusion_evaluated.write.format("delta").mode("overwrite").option("overwriteSchema", "true").saveAsTable(RESULT_AUDIT_TABLE))

# 全対象を除外した row だけを最終結果から取り除く。
# 除外の確認用列は RESULT_AUDIT_TABLE にだけ保持し、最終結果からは削除する。
result = (
    exclusion_evaluated
    .where(F.col("除外結果") != "全対象を除外")
    .drop("除外結果", "除外したドメイン", "一致した許可ドメイン")
)

(result.write.format("delta").mode("overwrite").option("overwriteSchema", "true").saveAsTable(TARGET_TABLE))

display(result.orderBy("submitted_datetime", "workspace_id", "notebook_id", "livy_id"))