# -*- coding: utf-8 -*-
import arcpy
import os
import re


def update_sde_connection_in_mxd(mxd_path, old_ip, new_ip):
    """
    更新 MXD 文件中所有图层的 SDE 连接 IP

    参数：
    mxd_path: MXD 文件路径
    old_ip: 旧的 IP 地址（可以是部分匹配）
    new_ip: 新的 IP 地址
    """
    try:
        # 打开 MXD
        mxd = arcpy.mapping.MapDocument(mxd_path)
        print("\n处理文件: {os.path.basename(mxd_path)}")

        # 更新数据框中的所有图层
        layers_updated = 0
        for df in arcpy.mapping.ListDataFrames(mxd):
            for lyr in arcpy.mapping.ListLayers(mxd, "", df):
                if lyr.supports("DATASOURCE"):
                    # 获取当前数据源
                    current_source = lyr.dataSource

                    # 检查是否为 SDE 连接并包含旧 IP
                    if "sde" in current_source.lower() and old_ip in current_source:
                        # 构建新的数据源路径
                        new_source = current_source.replace(old_ip, new_ip)

                        try:
                            # 更新图层数据源
                            lyr.replaceDataSource(
                                os.path.dirname(new_source),  # 工作空间路径
                                "SDE_WORKSPACE",  # 工作空间类型
                                os.path.basename(new_source),  # 数据集名称
                                False  # 不验证
                            )
                            layers_updated += 1
                            print("  已更新: {lyr.name}")
                            print("    原连接: {current_source}")
                            print("    新连接: {new_source}")
                        except Exception as e:
                            print("  更新失败 {lyr.name}: {str(e)}")

        # 保存 MXD
        if layers_updated > 0:
            mxd.save()
            print("  保存成功，更新了 {layers_updated} 个图层")
        else:
            print("  未找到需要更新的 SDE 连接")

        # 释放 MXD 对象
        del mxd

        return True

    except Exception as e:
        print("处理失败 {mxd_path}: {str(e)}")
        return False


def batch_update_sde_connections(folder_path, old_ip, new_ip, recursive=False):
    """
    批量处理文件夹中的 MXD 文件

    参数：
    folder_path: 文件夹路径
    old_ip: 旧的 IP 地址
    new_ip: 新的 IP 地址
    recursive: 是否递归处理子文件夹
    """
    # 支持的 MXD 文件扩展名
    extensions = ['.mxd', '.MXD']

    # 收集所有 MXD 文件
    mxd_files = []

    if recursive:
        for root, dirs, files in os.walk(folder_path):
            for file in files:
                if os.path.splitext(file)[1] in extensions:
                    mxd_files.append(os.path.join(root, file))
    else:
        for file in os.listdir(folder_path):
            if os.path.splitext(file)[1] in extensions:
                mxd_files.append(os.path.join(folder_path, file))

    print("找到 {len(mxd_files)} 个 MXD 文件")

    # 处理每个 MXD 文件
    success_count = 0
    for mxd_file in mxd_files:
        if update_sde_connection_in_mxd(mxd_file, old_ip, new_ip):
            success_count += 1

    print("\n处理完成！成功处理 {success_count}/{len(mxd_files)} 个文件")


def update_sde_connection_file(connection_file, old_ip, new_ip):
    """
    直接更新 SDE 连接文件 (.sde)
    """
    try:
        # 读取连接文件内容
        with open(connection_file, 'r') as f:
            content = f.read()

        # 替换 IP 地址
        new_content = content.replace(old_ip, new_ip)

        # 保存文件
        with open(connection_file, 'w') as f:
            f.write(new_content)

        print("已更新连接文件: {os.path.basename(connection_file)}")
        return True

    except Exception as e:
        print("更新连接文件失败: {str(e)}")
        return False


# ==================== 使用示例 ====================

if __name__ == "__main__":
    # 示例 1: 处理单个 MXD 文件
    print("=== 示例 1: 处理单个 MXD 文件 ===")
    # update_sde_connection_in_mxd(
    #     r"C:\Project\maps\project.mxd",
    #     "192.168.1.100",  # 旧 IP
    #     "10.255.133.176"  # 新 IP
    # )

    # 示例 2: 批量处理文件夹
    print("\n=== 示例 2: 批量处理文件夹 ===")
    batch_update_sde_connections(
        r"C:\TEMP\1",  # 文件夹路径
        "192.168.100.252",  # 旧 IP
        "192.168.60.20",  # 新 IP
        recursive=True  # 递归处理子文件夹
    )