"""
   @Author      : CT
   @Description : 解析全国邮编 json 文件（63万+条），写入 Hive 表
                  dim_china_zipcode（表由人工建好，本脚本只负责解析+写入）
                  文件从服务器本地读取（不走 HDFS）：用 Python 原生 json.load 在 driver 端
                  单机解析，解析完再用 spark.createDataFrame 分发成 DataFrame——这样避免了
                  spark.read 直接读本地路径时，任务被调度到其它 executor 节点、那台机器本地
                  没有这个文件导致 FileNotFoundException 的问题（file:// 读的是"任务实际
                  跑在哪台机器就读哪台机器本地磁盘"，driver 端 Python 原生读文件才是只读
                  提交任务这台机器本地磁盘，不受 executor 调度影响）
                  代价：63万+条数据在 driver 端单机解析，比 spark.read 分布式读要慢，但避免
                  了必须先把文件传到 HDFS 上的额外步骤
                  源 json 顶层直接是一个数组，每个元素字段：省、市、区、街道、村/居委会、邮编、
                  完整地址，7个字段全部入库
                  写入 Hive 前的处理：
                    - 空字符串统一转成 None（源数据里字段缺失是用 "" 表示的，不是 null）
                    - "市" 字段是"市辖区"的（北京/天津/上海/重庆这类直辖市），替换成对应的"省"
                      字段值（比如 省=北京市，市=市辖区 -> 市=北京市）
                  按如下顺序写入 Hive（insertInto 按列位置对应，务必确保建表时列顺序和下面一致）：
                    province(省)、city(市)、district(区)、street(街道)、village(村/居委会)、
                    zipcode(邮编)、full_address(完整地址)
                  写入前的校验：
                    7个字段完全相同的重复行，自动去重，并把去重前的重复记录打印出来，不报错
                    阻断流程（源数据里确实存在个别区县被重复列了好几遍的情况，去重不丢信息）
   @Usage       : python dim_china_zipcode.py <json_file_path>
   @CreateTime  : 2026-07-15
   @UpdateTime  : 2026-07-16
"""

import json
import os
import sys

sys.path.append(os.path.dirname(sys.path[0]))

from utils.common_util import CommonUtil
from utils.spark_util import SparkUtil
from pyspark.sql.types import StructType, StructField, StringType

HIVE_TABLE = "dim_china_zipcode"

DEFAULT_FILE_PATH = r"/home/chenyuanjie/china_postcode.json"

FIELD_MAP = [
    ("省", "province"),
    ("市", "city"),
    ("区", "district"),
    ("街道", "street"),
    ("村/居委会", "village"),
    ("邮编", "zipcode"),
    ("完整地址", "full_address"),
]

MUNICIPALITY_CITY_PLACEHOLDER = "市辖区"


def build_rows(data):
    rows = []
    for item in data:
        values = [item.get(cn_key) or None for cn_key, _ in FIELD_MAP]
        province, city = values[0], values[1]
        if city == MUNICIPALITY_CITY_PLACEHOLDER:
            values[1] = province
        rows.append(tuple(values))

    if not rows:
        raise ValueError("源数据是空数组，请检查文件内容")

    print(f"共解析 {len(rows)} 条记录（去重前）")
    return rows


def dedupe_rows(rows):
    seen = {}
    duplicates = []
    deduped = []
    for row in rows:
        if row in seen:
            duplicates.append(row)
            continue
        seen[row] = True
        deduped.append(row)

    if duplicates:
        print(f"发现 {len(duplicates)} 行完全重复（7个字段都一样），已去重，重复的记录如下：")
        for row in duplicates:
            print(f"  {row}")

    return deduped


class DimChinaZipcode(object):

    def __init__(self, json_path):
        self.json_path = json_path
        self.spark = SparkUtil.get_spark_session(f"{self.__class__.__name__}")
        self.df_save = None
        self.raw_total = None

    def read_data(self):
        print(f"读取本地文件：{self.json_path}")
        with open(self.json_path, encoding="utf-8") as f:
            data = json.load(f)

        rows = build_rows(data)
        self.raw_total = len(rows)
        rows = dedupe_rows(rows)

        schema = StructType([
            StructField(en_key, StringType(), True) for _, en_key in FIELD_MAP
        ])
        # 后面 validate_data/save_data 要对同一份数据跑好几次 count/dropDuplicates，
        # 不 cache 的话每次 action 都会重新触发一遍 createDataFrame 的分发
        self.df_save = self.spark.createDataFrame(rows, schema=schema).cache()

    def validate_data(self):
        total = self.df_save.count()
        distinct_total = self.df_save.dropDuplicates().count()
        if distinct_total != total:
            # 正常不会走到这，dedupe_rows 已经去过重了，留着是防止后续逻辑改动引入新的重复
            raise ValueError(
                f"去重后仍存在完全重复的行：总行数 {total}，dropDuplicates 后 {distinct_total}"
            )

        print(f"校验通过：去重前 {self.raw_total} 条，去重后 {total} 条且7个字段组合唯一")

    def save_data(self):
        total = self.df_save.count()
        print(f"写入 Hive 表 {HIVE_TABLE}，共 {total} 条")
        CommonUtil.save_or_update_table(
            spark_session=self.spark,
            hive_tb_name=HIVE_TABLE,
            partition_dict={},
            df_save=self.df_save,
        )
        print("写入完毕")

    def run(self):
        self.read_data()
        self.validate_data()
        self.save_data()


if __name__ == "__main__":
    json_path = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_FILE_PATH

    if not os.path.exists(json_path):
        print(f"文件不存在：{json_path}")
        sys.exit(1)

    DimChinaZipcode(json_path=json_path).run()
