import os
import json
import base64
import requests
import click
import boto3
from botocore.exceptions import ClientError
from tqdm import tqdm
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad

class ZsyClient:
    def __init__(self, debug=False):
        self.debug = debug
        self.center_host = os.getenv("CENTER_HOST", "cloud.ai4net.cn")
        self.zone = os.getenv("ZONE")
        if not self.zone:
            service_prefix = os.getenv("SERVICE_PREFIX")
            if service_prefix:
                self.zone = service_prefix.split('/')[0]
        
        if not self.zone:
            raise click.ClickException("错误: 未能通过环境变量 ZONE 或 SERVICE_PREFIX 获取到有效的区域信息 (ZONE)")
        
        # 共享密钥 (需为 16, 24 或 32 字节)
        self.aes_key = b"r&%Njt#wam5lVVj$" 

    @staticmethod
    def aes_encrypt(key, data_dict):
        """AES-CBC 加密 (模拟后端)"""
        data_str = json.dumps(data_dict)
        cipher = AES.new(key, AES.MODE_CBC)
        padded_data = pad(data_str.encode('utf-8'), AES.block_size)
        iv = cipher.iv
        ciphertext = cipher.encrypt(padded_data)
        # 将 IV 和 密文 拼接并 Base64 编码
        return base64.b64encode(iv + ciphertext).decode('utf-8')

    @staticmethod
    def aes_decrypt(key, encrypted_str):
        """AES-CBC 解密 (客户端使用)"""
        raw_data = base64.b64decode(encrypted_str)
        iv = raw_data[:AES.block_size]
        ciphertext = raw_data[AES.block_size:]
        cipher = AES.new(key, AES.MODE_CBC, iv)
        padded_data = cipher.decrypt(ciphertext)
        data_str = unpad(padded_data, AES.block_size).decode('utf-8')
        return json.loads(data_str)

    def fetch_auth_data(self, resource_type, resource_id):
        """调用远程接口获取认证数据"""
        # 使用 http 还是 https 取决于环境，这里默认优先从环境变量获取完整 host
        # 如果 host 包含端口或不包含域名，可能需要调整协议
        protocol = "https" if "localhost" not in self.center_host else "http"
        url = f"{protocol}://{self.center_host}/center/server/v1/zsyctl/auth"
        
        params = {
            "type": resource_type,
            "id": resource_id,
            "zone": self.zone
        }
        
        try:
            response = requests.get(url, params=params, timeout=10)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            raise click.ClickException(f"获取认证信息失败")

    def get_auth_token(self, resource_type, resource_id):
        """获取加密后的 Token/String"""
        data = self.fetch_auth_data(resource_type, resource_id)
        return data.get("token")

    def decode_token(self, token):
        """解密 AES 字符串获取 S3 凭证"""
        try:
            config = self.aes_decrypt(self.aes_key, token)
            
            # 如果后端没返回地址，则根据 center_host 自己拼凑 (cloud.ai4net.cn -> ai4net.cn)
            domain = self.center_host.split('.')
            if len(domain) >= 2:
                base_domain = '.'.join(domain[-2:])
                # 去掉可能存在的端口号 (例如 localhost:8080 -> localhost)
                base_domain = base_domain.split(':')[0]
            else:
                base_domain = self.center_host.split(':')[0]

            if not config.get("centerS3Url"):
                config["centerS3Url"] = f"https://maas-zsyfs.{base_domain}"
            
            if not config.get("currentS3Url"):
                # 如果是 localhost 测试，通常也指向中心或者特定的测试 S3
                if "localhost" in self.center_host:
                    config["currentS3Url"] = config["centerS3Url"]
                else:
                    config["currentS3Url"] = f"https://{self.zone}-maas-zsyfs.{base_domain}"
            
            if self.debug:
                ak = config.get('accessKey', '')
                sk = config.get('secretKey', '')
            
            return config
        except Exception as e:
            raise click.ClickException(f"解密凭证失败 (请检查密钥): {str(e)}")

    def download_resource(self, s3_config, local_dir):
        """执行 S3 下载 (支持目录递归下载 + 进度条)"""
        access_key = s3_config.get("accessKey")
        secret_key = s3_config.get("secretKey")
        bucket = s3_config.get("bucket")
        prefix = s3_config.get("key")
        
        # 补全目录前缀的斜杠
        if not prefix.endswith('/'):
            prefix += '/'

        # 优先级：先 current，后 center
        endpoints = [s3_config.get("currentS3Url"), s3_config.get("centerS3Url")]
        
        for endpoint in endpoints:
            if not endpoint:
                continue
            
            try:
                s3 = boto3.client(
                    's3',
                    aws_access_key_id=access_key,
                    aws_secret_access_key=secret_key,
                    endpoint_url=endpoint,
                    config=boto3.session.Config(signature_version='s3v4')
                )
                
                # 列出前缀下的所有对象
                paginator = s3.get_paginator('list_objects_v2')
                pages = list(paginator.paginate(Bucket=bucket, Prefix=prefix))
                
                # 计算总大小和文件数
                all_objects = []
                total_size = 0
                for page in pages:
                    if 'Contents' not in page:
                        continue
                    for obj in page['Contents']:
                        if not obj['Key'].endswith('/'):
                            all_objects.append(obj)
                            total_size += obj['Size']
                
                # 使用 tqdm 显示总体进度
                if not all_objects:
                    click.echo(f"在 {endpoint} 未找到对象")
                    continue

                with tqdm(total=total_size, unit='B', unit_scale=True, desc="总体进度") as pbar:
                    for obj in all_objects:
                        key = obj['Key']
                        relative_path = os.path.relpath(key, prefix)
                        local_file_path = os.path.join(local_dir, relative_path)
                        
                        os.makedirs(os.path.dirname(local_file_path), exist_ok=True)
                        
                        # 下载并更新进度条
                        s3.download_file(
                            bucket, 
                            key, 
                            local_file_path,
                            Callback=lambda bytes_transferred: pbar.update(bytes_transferred)
                        )
                
                click.echo(f"\n下载完成！文件已保存至: {local_dir}")
                return
            
            except ClientError as e:
                error_code = e.response.get('Error', {}).get('Code', 'Unknown')
                if error_code in ['SignatureDoesNotMatch', 'InvalidAccessKeyId', 'AccessDenied']:
                    raise click.ClickException(f"授权错误: S3 凭证校验失败 ({error_code})。请检查 SecretKey 是否配置正确。")
                
                if self.debug:
                    click.echo(f"  --> 从 {endpoint} 下载过程中出错: {error_code}: {str(e)}")
                continue
            except Exception as e:
                if self.debug:
                    click.echo(f"  --> 从 {endpoint} 下载过程中出错: {type(e).__name__}: {str(e)}")
                continue
        
        raise click.ClickException("资源下载失败：当前活跃集群和中心集群均无法连接或未找到文件。")

@click.group()
@click.option('--debug', is_flag=True, help='开启调试输出')
@click.pass_context
def cli(ctx, debug):
    """zsyctl - 智算云命令行工具"""
    ctx.ensure_object(dict)
    ctx.obj['DEBUG'] = debug

@cli.command()
@click.option('--dataset', 'dataset_id', help='数据集 ID')
@click.option('--model', 'model_id', help='模型 ID')
@click.option('--local_dir', required=True, help='本地输出路径')
@click.pass_context
def download(ctx, dataset_id, model_id, local_dir):
    """下载资源"""
    resource_id = dataset_id or model_id
    if not resource_id:
        raise click.UsageError("必须提供 --dataset 或 --model 参数")
    
    # 确定资源类型
    resource_type = "dataset" if dataset_id else "model"
    
    debug = ctx.obj.get('DEBUG', False)
    client = ZsyClient(debug=debug)
    
    click.echo(f"正在认证资源 {resource_type}: {resource_id}...")
    token = client.get_auth_token(resource_type, resource_id)
    
    s3_config = client.decode_token(token)
    
    # 执行目录递归下载
    client.download_resource(s3_config, local_dir)

def main():
    cli()

if __name__ == "__main__":
    main()
