import yaml
import oss2
from tqdm import tqdm
import os
import numpy as np
import os
import argparse
# import bpy
import sys
import open3d as o3d

from pathlib import Path

import requests
import json
import shutil

from config import url_get_oss_suffix_by_orderId

from general import transform_save_bpy
from general import read_models_from_json

from compute_print_net import get_oss_client

def download_transform_save_by_json(json_name, workdir, oss_config):
    layout_data = download_datas_by_json(json_name, workdir, oss_config)
    original_obj_pid_dir = workdir
    transform_save_bpy(layout_data, original_obj_pid_dir)

class DataTransfer:
    '''
    数据传输类
    '''
    def __init__(self, local_path: str, oss_path: str, oss_client: oss2.Bucket):
        '''
        local_path: 本地输出路径
        oss_path: oss路径
        oss_client: oss客户端
        '''
        self.local_path = local_path
        self.oss_path = oss_path.lstrip('/')
        self.oss_client = oss_client
        # self.description = description

    # @log_execution(self.description)
    def download_data(self):
        """
        从 OSS 下载数据到本地，保持原有目录结构
        """
        
        # 列出所有对象
        objects = []
        prefix = self.oss_path.lstrip('/')  # 移除开头的 '/' 以匹配 OSS 格式
        
        for obj in oss2.ObjectIterator(self.oss_client, prefix=prefix):
            if obj.key != prefix:  # 跳过目录本身
                objects.append(obj.key)
        
        # 下载所有文件，添加进度条
        for obj_key in tqdm(objects, desc="下载进度"):
            if obj_key.endswith('/'):
                continue

            if "printId" in obj_key:
                continue

            # 计算相对路径
            rel_path = obj_key[len(prefix):].lstrip('/')
            # 构建本地完整路径
            local_path = os.path.join(self.local_path, rel_path)

            # 创建必要的目录
            os.makedirs(os.path.dirname(local_path), exist_ok=True)
            
            # 下载文件
            self.oss_client.get_object_to_file(obj_key, local_path)
            
            print("download_data local_path=" + local_path)

    order_id: str
    pid: str
    model_height: str

    def download_data_rename_json(self, json_model_info):
        """
        从 OSS 下载数据到本地，保持原有目录结构
        """
        
        # 列出所有对象
        objects = []
        prefix = self.oss_path.lstrip('/')  # 移除开头的 '/' 以匹配 OSS 格式
        
        for obj in oss2.ObjectIterator(self.oss_client, prefix=prefix):
            if obj.key != prefix:  # 跳过目录本身
                objects.append(obj.key)
        
        # 下载所有文件，添加进度条
        for obj_key in tqdm(objects, desc="下载进度"):
            if obj_key.endswith('/'):
                continue

            if "printId" in obj_key:
                continue

            # 计算相对路径
            rel_path = obj_key[len(prefix):].lstrip('/')

            file_dir, file_name = os.path.split(rel_path)
            file_base, file_ext = os.path.splitext(file_name)

            # 根据文件后缀名进行重命名
            if file_ext.lower() in ['.mtl', '.jpg', '.jpeg', '.png']:
                # 对于.mtl和图片文件，在原名前加order_id
                new_file_name = f"{json_model_info.order_id}_{file_name}"
                # new_file_name = file_name
            elif file_ext.lower() == '.obj':
                # 对于.obj文件，完全重命名
                new_file_name = f"{json_model_info.obj_name}"
            else:
                # 其他文件类型保持原名
                new_file_name = file_name
            print("new_file_name=", new_file_name)

            # 构建新的相对路径
            if file_dir:  # 如果有子目录
                new_rel_path = os.path.join(file_dir, new_file_name)
            else:
                new_rel_path = new_file_name

            # 构建本地完整路径
            local_path = os.path.join(self.local_path, new_rel_path)
            
            # 创建必要的目录
            os.makedirs(os.path.dirname(local_path), exist_ok=True)
            
            # 下载文件
            self.oss_client.get_object_to_file(obj_key, local_path)

            if file_ext == '.obj':  # 10MB以上
                try:
                    # 使用临时文件避免内存问题 [8](@ref)
                    temp_path = local_path + '.tmp'
                    with open(local_path, 'r', encoding='utf-8') as f_in, \
                        open(temp_path, 'w', encoding='utf-8') as f_out:
                        
                        mtllib_modified = False
                        for line in f_in:
                            if not mtllib_modified and line.strip().startswith('mtllib '):
                                parts = line.split(' ', 1)
                                if len(parts) > 1:
                                    old_mtl_name = parts[1].strip()
                                    new_mtl_name = f"{json_model_info.order_id}_{old_mtl_name}"
                                    f_out.write(f"mtllib {new_mtl_name}\n")
                                    mtllib_modified = True
                                    continue
                            f_out.write(line)
                    
                    os.replace(temp_path, local_path)  # 原子性替换
                    
                except IOError as e:
                    print(f"处理大文件 {local_path} 时出错: {e}")
                    if os.path.exists(temp_path):
                        os.remove(temp_path)
            
            # 优化后的.obj文件处理逻辑
            if file_ext == '.mtl':
                try:
                    # 使用更高效的文件读取方式 [6,8](@ref)
                    with open(local_path, 'r', encoding='utf-8') as f:
                        content = f.read()
                    
                    # 使用字符串方法直接查找和替换，避免不必要的循环 [9](@ref)
                    lines = content.split('\n')
                    mtllib_modified = False
                    
                    for i, line in enumerate(lines):
                        stripped_line = line.strip()
                        if not mtllib_modified and stripped_line.startswith('map_Kd '):
                            # 更高效的分割方式 [9](@ref)
                            parts = line.split(' ', 1)
                            if len(parts) > 1:
                                old_name = parts[1].strip()
                                new_name = f"{json_model_info.order_id}_{old_name}"
                                lines[i] = f"map_Kd {new_name}"
                                mtllib_modified = True
                                # print(f"已更新材质库引用: {old_name} -> {new_name}")
                                break  # 找到第一个后立即退出
                    
                    # 批量写入，减少I/O操作 [6](@ref)
                    with open(local_path, 'w', encoding='utf-8') as f:
                        f.write('\n'.join(lines))
                        
                except IOError as e:
                    print(f"处理文件 {local_path} 时出错: {e}")
                except UnicodeDecodeError as e:
                    print(f"文件编码错误 {local_path}: {e}")

            print(f"下载文件: {obj_key} -> {local_path}")

    def download_data_rename_batch(self, batch_model_info):
        """
        从 OSS 下载数据到本地，保持原有目录结构
        """
        
        # 列出所有对象
        objects = []
        prefix = self.oss_path.lstrip('/')  # 移除开头的 '/' 以匹配 OSS 格式

        prefix_exists = False
        
        for obj in oss2.ObjectIterator(self.oss_client, prefix=prefix):
            prefix_exists = True
            if obj.key != prefix:  # 跳过目录本身
                objects.append(obj.key)
                # print(f"obj.key={obj.key}")
        
        if not prefix_exists:
            print(f"前缀 '{prefix}' 下没有找到任何文件或目录。")
            return False
        else:
            print(f"前缀 '{prefix}' 存在，共找到 {len(objects)} 个对象。")

        # 下载所有文件，添加进度条
        for obj_key in tqdm(objects, desc="下载进度"):
            if obj_key.endswith('/'):
                # print("下载 endswith('/'")
                continue

            if "printId" in obj_key:
                # print(f"下载 in obj_key")
                continue

            # 计算相对路径
            rel_path = obj_key[len(prefix):].lstrip('/')

            file_dir, file_name = os.path.split(rel_path)
            file_base, file_ext = os.path.splitext(file_name)

            # 根据文件后缀名进行重命名
            if file_ext.lower() in ['.mtl', '.jpg', '.jpeg', '.png']:
                # 对于.mtl和图片文件，在原名前加order_id
                new_file_name = f"{batch_model_info.order_id}_{file_name}"
                # new_file_name = file_name
            elif file_ext.lower() == '.obj':
                # 对于.obj文件，完全重命名
                new_file_name = f"{batch_model_info.order_id}_{batch_model_info.pid}_P{batch_model_info.print_order_id}_{batch_model_info.model_size}{file_ext}"
            else:
                # 其他文件类型保持原名
                new_file_name = file_name

            # 构建新的相对路径
            if file_dir:  # 如果有子目录
                new_rel_path = os.path.join(file_dir, new_file_name)
            else:
                new_rel_path = new_file_name

            # 构建本地完整路径
            local_path = os.path.join(self.local_path, new_rel_path)
            
            # 创建必要的目录
            os.makedirs(os.path.dirname(local_path), exist_ok=True)
            
            # 下载文件
            self.oss_client.get_object_to_file(obj_key, local_path)

            if file_ext == '.obj':  # 10MB以上
                try:
                    # 使用临时文件避免内存问题 [8](@ref)
                    temp_path = local_path + '.tmp'
                    with open(local_path, 'r', encoding='utf-8') as f_in, \
                        open(temp_path, 'w', encoding='utf-8') as f_out:
                        
                        mtllib_modified = False
                        for line in f_in:
                            if not mtllib_modified and line.strip().startswith('mtllib '):
                                parts = line.split(' ', 1)
                                if len(parts) > 1:
                                    old_mtl_name = parts[1].strip()
                                    new_mtl_name = f"{batch_model_info.order_id}_{old_mtl_name}"
                                    f_out.write(f"mtllib {new_mtl_name}\n")
                                    mtllib_modified = True
                                    print("len(parts) > 1")
                                    continue
                            f_out.write(line)
                    
                    os.replace(temp_path, local_path)  # 原子性替换
                    
                except IOError as e:
                    print(f"处理大文件 {local_path} 时出错: {e}")
                    if os.path.exists(temp_path):
                        os.remove(temp_path)
            
            # 优化后的.obj文件处理逻辑
            if file_ext == '.mtl':
                try:
                    # 使用更高效的文件读取方式 [6,8](@ref)
                    with open(local_path, 'r', encoding='utf-8') as f:
                        content = f.read()
                    
                    # 使用字符串方法直接查找和替换，避免不必要的循环 [9](@ref)
                    lines = content.split('\n')
                    mtllib_modified = False
                    
                    for i, line in enumerate(lines):
                        stripped_line = line.strip()
                        if not mtllib_modified and stripped_line.startswith('map_Kd '):
                            # 更高效的分割方式 [9](@ref)
                            parts = line.split(' ', 1)
                            if len(parts) > 1:
                                old_name = parts[1].strip()
                                new_name = f"{batch_model_info.order_id}_{old_name}"
                                lines[i] = f"map_Kd {new_name}"
                                mtllib_modified = True
                                # print(f"已更新材质库引用: {old_name} -> {new_name}")
                                break  # 找到第一个后立即退出
                    
                    # 批量写入，减少I/O操作 [6](@ref)
                    with open(local_path, 'w', encoding='utf-8') as f:
                        f.write('\n'.join(lines))
                        
                except IOError as e:
                    print(f"处理文件 {local_path} 时出错: {e}")
                except UnicodeDecodeError as e:
                    print(f"文件编码错误 {local_path}: {e}")

            print(f"下载文件: {obj_key} -> {local_path}")
        
        return True

    def download_single_file(self):
        """
        下载单个文件从OSS到本地
        """
        # 确保本地目录存在
        os.makedirs(os.path.dirname(self.local_path), exist_ok=True)
        
        # 直接下载文件
        try:
            self.oss_client.get_object_to_file(self.oss_path, self.local_path)
            print(f"文件已下载到: {self.local_path}")
        except oss2.exceptions.NoSuchKey:
            print(f"OSS文件不存在: {self.oss_path}")

    def upload_data(self):
        '''
        上传数据到OSS
        '''
        # 检测本地路径是否存在
        if not os.path.exists(self.local_path):
            raise FileNotFoundError(f"本地路径不存在: {self.local_path}")

        # 判断本地路径是文件还是目录
        if os.path.isfile(self.local_path):
            local_suffix = Path(self.local_path).suffix
            oss_suffix = Path(self.oss_path).suffix

            if oss_suffix and oss_suffix != local_suffix:
                # 后缀名不一致，上传到指定文件夹下的同名文件
                oss_dir = os.path.dirname(self.oss_path)
                oss_target_path = os.path.join(oss_dir, os.path.basename(self.local_path))
            else:
                # 后缀名一致，上传到指定OSS路径
                oss_target_path = self.oss_path

            # 上传文件
            self.oss_client.put_object_from_file(oss_target_path, self.local_path)

            print(f"文件已上传到: {oss_target_path}")

        elif os.path.isdir(self.local_path):
            oss_suffix = Path(self.oss_path).suffix
            if oss_suffix:
                raise ValueError("不能将目录上传到具有后缀名的OSS路径。")

            # 遍历本地目录并上传
            for root, dirs, files in os.walk(self.local_path):
                for file in files:
                    local_file_path = os.path.join(root, file)
                    relative_path = os.path.relpath(local_file_path, self.local_path)
                    oss_file_path = os.path.join(self.oss_path, relative_path).replace("\\", "/")
                    
                    # 创建必要的目录
                    oss_dir = os.path.dirname(oss_file_path)
                    
                    # 上传文件
                    self.oss_client.put_object_from_file(oss_file_path, local_file_path)
                    print(f"文件已上传到: {oss_file_path}")
        else:
            raise ValueError(f"无效的本地路径类型: {self.local_path}")

def get_api(url):
    try:
        response = requests.get(url)
        response.raise_for_status()  # 检查请求是否成功
        response = json.loads(response.text)
        if response.get("code") != 1000:
            raise Exception(f"Error fetching URL {url}: {response.get('message')}")
        else:
            return response
    except requests.exceptions.RequestException as e:
        raise Exception(f"Error fetching URL {url}: {e}")
    
def download_data_by_json(model_info, workdir, oss_client ):
    try:
        pid = model_info.pid
        model_height = model_info.model_height
        target_dir = f"{workdir}"

        url = f"{url_get_oss_suffix_by_orderId}{model_info.order_id}"
        res = requests.get(url)

        data = res.json()["data"]
        # print("datas=",data)
        data = data.replace("/init_obj", "")

        print("target_dir=", target_dir)
        download_textures = DataTransfer(target_dir, f"objs/download/print/{pid}/{data}/{model_height}/", oss_client)

        download_textures.download_data_rename_json(model_info)
        # 下载后检查目标文件夹是否为空
        if os.path.exists(target_dir) and not os.listdir(target_dir):
            shutil.rmtree(target_dir)
            print(f"下载后检查发现目标文件夹为空，已删除: {target_dir}")
    except Exception as e:
        print(f"下载失败: {pid}, 错误: {str(e)}")
        pass

def download_datas_by_json(pid_file, workdir, oss_config):
    oss_client = get_oss_client(oss_config)

    json_path = os.path.join(workdir, f"{pid_file}.json")
    
    # 读取所有PID
    list_model_info, json_data = read_models_from_json(json_path)
    print(f"从文件读取了 {len(list_model_info)} 个PID")
    # 批量下载
    for model_info in list_model_info:
        print(f"开始下载PID: {model_info}")
        download_data_by_json(model_info, args.workdir, oss_client)

    return json_data

if __name__ == "__main__":

    script_args = []

    try:
        separator_index = sys.argv.index("--")
        script_args = sys.argv[separator_index + 1:]
    except ValueError:
        pass

    parser = argparse.ArgumentParser()

    parser.add_argument("--pid_file",  type=str, required=True, help="批次号, 也是json文件名")
    parser.add_argument("--workdir", type=str, required=True, help="json文件所在的目录")
    parser.add_argument("--oss_config", type=str, required=True, help="run.yaml所在的目录")
    args = parser.parse_args(script_args)

    download_transform_save_by_json(args.pid_file, args.workdir, args.oss_config)

    # blender --background --python download_print_out.py -- --pid_file {your_batch_id} --workdir {your_batch_dir} --oss_config {your_yaml_dir/run.yaml}