import json
import os
import urllib

import loguru
import yaml
from qcloud_cos import CosConfig, CosS3Client, CosClientError, CosServiceError


class COSDownload:
    """腾讯云COS下载"""

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "cos_bucket": ("STRING", {"default": "bwkj-cos-1324682537"}),
                "cos_key": ("STRING", {"multiline": True}),
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("视频存储路径",)
    FUNCTION = "download"
    CATEGORY = "不忘科技-自定义节点🚩"

    def download(self, cos_bucket, cos_key):
        loguru.logger.info("download to {}".format(os.path.join(
                    os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                    "download",
                    os.path.dirname(cos_key),
                )))
        if os.sep in cos_key or "/" in cos_key or "\\" in cos_key:
            os.makedirs(
                os.path.join(
                    os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                    "download",
                    os.path.dirname(cos_key),
                ),
                exist_ok=True,
            )
        for i in range(0, 10):
            try:
                with open(
                        os.path.join(
                            os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "config.yaml"
                        ),
                        encoding="utf-8",
                        mode="r+",
                ) as f:
                    yaml_config = yaml.load(f, Loader=yaml.FullLoader)
                config = CosConfig(
                    Region=yaml_config["region"],
                    SecretId=yaml_config["secret_id"],
                    SecretKey=yaml_config["secret_key"],
                )
                client = CosS3Client(config)
                response = client.download_file(
                    Bucket=cos_bucket,
                    Key=cos_key,
                    DestFilePath=os.path.join(
                        os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                        "download",
                        os.path.dirname(cos_key),
                        os.path.basename(cos_key),
                    ),
                )
                break
            except CosClientError or CosServiceError as e:
                print(f"下载失败 {e}")
        return (
            os.path.join(
                os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                "download",
                os.path.dirname(cos_key),
                os.path.basename(cos_key),
            ),
        )


class COSUpload:
    """腾讯云COS上传"""

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "cos_bucket": ("STRING", {"default": "bwkj-cos-1324682537"}),
                "path": ("STRING", {"multiline": True}),
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("COS文件Key",)

    FUNCTION = "upload"
    CATEGORY = "不忘科技-自定义节点🚩"

    def upload(self, cos_bucket, path):
        for i in range(0, 10):
            try:
                with open(
                        os.path.join(
                            os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "config.yaml"
                        ),
                        encoding="utf-8",
                        mode="r+",
                ) as f:
                    yaml_config = yaml.load(f, Loader=yaml.FullLoader)
                config = CosConfig(
                    Region=yaml_config["region"],
                    SecretId=yaml_config["secret_id"],
                    SecretKey=yaml_config["secret_key"],
                )
                client = CosS3Client(config)
                response = client.upload_file(
                    Bucket=cos_bucket,
                    Key="/".join(
                        [
                            yaml_config["subfolder"],
                            (
                                path.split("/")[-1]
                                if "/" in path
                                else path.split("\\")[-1]
                            ),
                        ]
                    ),
                    LocalFilePath=path,
                )
                break
            except CosClientError or CosServiceError as e:
                raise RuntimeError("上传失败")
        data = {"prompt_id": "",
                "video_url": "https://{}.cos.{}.myqcloud.com/{}".format(cos_bucket, yaml_config['region'],
                                                                        '/'.join([yaml_config['subfolder'],
                                                                                  path.split('/')[
                                                                                      -1] if '/' in path else
                                                                                  path.split('\\')[-1], ]))
                }
        headers = {'Content-Type': 'application/json'}
        try:
            req = urllib.request.Request("", data=json.dumps(data).encode("utf-8"), headers=headers)
            response = urllib.request.urlopen(req)
        except:
            raise RuntimeError("上报MQ状态失败")
        return (
            "/".join(
                [
                    yaml_config["subfolder"],
                    path.split("/")[-1] if "/" in path else path.split("\\")[-1],
                ]
            ),
        )