from lib.all import *
import subprocess
import json

host = "https://wishpal-ironfan-media.oss-cn-hangzhou-internal.aliyuncs.com/"


def get_pixel_format(video_url):
    command = [
        'ffprobe', '-v', 'quiet',
        '-show_streams', '-select_streams', 'v:0',
        '-print_format', 'json', video_url
    ]

    result = subprocess.run(command, capture_output=True, text=True)
    output = result.stdout

    data = json.loads(output)
    streams = data.get('streams', [])

    if streams:
        pix_fmt = streams[0].get('pix_fmt', 'Unknown')
        return pix_fmt
    else:
        return 'No video stream found'


class S:
    def __init__(self):
        self.col_video = MongoDB(
            host="mongodb://root:Wishpal2024@dds-bp1da1ddd62bede41.mongodb.rds.aliyuncs.com:3717,dds-bp1da1ddd62bede42.mongodb.rds.aliyuncs.com:3717/admin?replicaSet=mgset-77304659",
            port=3717,
            db="media",
            collection="video"
        )

    def __del__(self):
        self.col_video.close()

    def update_pixel_format(self, vid, pixel_format):
        q = {"_id": vid}
        up = {
            "$set": {
                "pixel_format": pixel_format
            }
        }
        ret = self.col_video.update_one(q, up)
        return ret

    def proc(self):
        prj = {
            "src_id": 1,
            "src_id_720": 1
        }
        docs = self.col_video.find({"pixel_format": {"$exists": False}}, projection=prj)

        total = len(docs)
        idx = 0
        for d in docs:
            idx += 1
            vid = safe_get_int(d, "_id")
            src_id_720 = safe_get_str(d, "src_id_720")
            url = host + src_id_720
            if len(src_id_720) > 0:
                # print(vid, url)
                pixel_format = get_pixel_format(url)
                ret = self.update_pixel_format(vid, pixel_format)
                print("{}/{} vid: {}, pf: {}, ret: {}".format(idx, total, vid, pixel_format, ret.modified_count))
                # if pixel_format != "yuv420p":
                #    print(vid, pixel_format)


# video_url = 'https://filecdn01.tiefen.fun/vdprodh264720/5a/d7/175d-2137-4a66-8e5c-d5ed61241b6c'
# pixel_format = get_pixel_format(video_url)
# print(f'Pixel Format: {pixel_format}')
s = S()
s.proc()