import os
import sys
import paramiko
from datetime import datetime
import time

def sftp_transfer():
    # 获取当前日期
    today_date = datetime.now().strftime('%Y%m%d')

    # sftp服务器
    host = '220.203.233.97'
    port = 16888
    username = 'sysop'
    password = 'swjg_Boyizc1212'

    # 远程路径
    remote_base_path = f'/home/sysop/kr_md_save_centos7_9_v3/data/stockdata/{today_date}'

    # 本地路径
    local_base_path = f'G:\\data_new\\stockdata\\{today_date}'

    # 如果本地目录不存在，则创建
    if not os.path.exists(local_base_path):
        os.makedirs(local_base_path, exist_ok=True)

    # 连接SFTP服务器
    try:
        transport = paramiko.Transport((host, port))
        transport.connect(username=username, password=password)
        sftp = paramiko.SFTPClient.from_transport(transport)

        print(f"成功连接到SFTP服务器: {host}:{port}")
        print(f"远程目录: {remote_base_path}")
        print(f"本地目录: {local_base_path}")

        # 列出远程目录中的所有文件
        try:
            remote_files = sftp.listdir(remote_base_path)
        except FileNotFoundError:
            print(f"错误: 远程目录不存在: {remote_base_path}")
            print("可能原因: 1) 日期格式错误 2) 远程服务器上该日期的数据还未生成")
            return False

        if not remote_files:
            print("警告: 远程目录为空，没有文件需要传输")
            return True

        print(f"找到 {len(remote_files)} 个文件需要传输")

        # 传输文件
        transferred_count = 0
        for filename in remote_files:
            remote_file_path = os.path.join(remote_base_path, filename)
            local_file_path = os.path.join(local_base_path, filename)

            try:
                print(f"正在传输: {filename} ...")
                sftp.get(remote_file_path, local_file_path)
                transferred_count += 1
                print(f"  完成")
            except Exception as e:
                print(f"  传输失败: {str(e)}")

        print(f"传输完成: 成功传输 {transferred_count}/{len(remote_files)} 个文件")

        # 关闭连接
        sftp.close()
        transport.close()

        return transferred_count > 0

    except Exception as e:
        print(f"连接或传输过程中发生错误: {str(e)}")
        return False


def main():
    print("=" * 60)
    print(f"开始执行SFTP数据同步任务")
    print(f"时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print("=" * 60)

    # 记录开始时间
    start_time = time.time()

    # 执行传输任务
    success = sftp_transfer()

    # 记录结束时间
    end_time = time.time()
    elapsed_time = end_time - start_time

    print("=" * 60)
    if success:
        print(f"任务执行成功! 耗时: {elapsed_time:.2f}秒")
    else:
        print(f"任务执行失败! 耗时: {elapsed_time:.2f}秒")
    print("=" * 60)

    return 0 if success else 1


if __name__ == '__main__':
    sys.exit(main())