import os
import json
import hashlib
import pandas as pd # Import pandas for spreadsheet reading

# Constants for base URLs
BASE_URL_5GVIDEO = "https://dash-large-files.akamaized.net/WAVE/3GPP/5GVideo/ReferenceSequences/"
BASE_URL_CONFORMANCE = "https://dash-large-files.akamaized.net/WAVE/3GPP/Conformance/26265/Development/Test%20Sequences/"

def add_component(data_dict, *, seq_type, comp_type, hero, uri, md5, size):
    """
    向 data_dict['Signal']['Components'] 追加一个分量。
    - 当 data_dict['Signal']['type'] == 'stereo' 时，分量对象含 'type'、'hero'（且排在最前面）
    - 其他类型（如 'mono'）不包含 'type'、'hero'
    """
    if "Signal" not in data_dict or "Components" not in data_dict["Signal"]:
        raise KeyError("data_dict['Signal']['Components'] not exist")

    if seq_type == "stereo":
        comp = {
            "type": comp_type,
            "hero": hero,
            "URI": uri,
            "md5": md5,
            "size": size
        }
    else:
        # no 'type' & 'hero' when not stereo
        comp = {
            "URI": uri,
            "md5": md5,
            "size": size
        }
    data_dict["Signal"]["Components"].append(comp)

def _yuv_bytes_factor(subsampling: str) -> float:
    """
    根据采样格式返回每像素通道系数（YUV 裸数据的字节估算用）
    """
    ss = (subsampling or "").strip()
    if ss == "420":
        return 1.5
    if ss == "422":
        return 2.0
    if ss == "444":
        return 3.0
    # default 420
    return 1.5

def get_md5_value(file_path, md5_from_csv):
    """
    Prioritizes MD5 from CSV, falls back to local calculation if not provided.
    """
    # Check for pandas NaN or empty string
    if pd.isna(md5_from_csv) or str(md5_from_csv).strip() == '':
        if os.path.isfile(file_path):
            md5_hash = hashlib.md5()
            try:
                with open(file_path, "rb") as f:
                    for byte_block in iter(lambda: f.read(4096), b""):
                        md5_hash.update(byte_block)
                return md5_hash.hexdigest()
            except FileNotFoundError:
                return "" # Handle case where file might be specified but not exist
        return "" # Return empty string if file doesn't exist and no MD5 from CSV
    return str(md5_from_csv).strip()

def get_datasize(row_data, yuv_name, row_index):
    """
    计算并返回：(datasize, bytesize)
      - datasize: 实际文件大小（字节），若文件不存在则为 None
      - bytesize: 按 width/height/subsampling/bitDepth/frameCount 估算的期望大小（字节）
    同时打印一致性信息（missing / ok / mismatch）。
    """
    exists = os.path.isfile(yuv_name)
    datasize = os.path.getsize(yuv_name) if exists else None

    # Accessing data using iloc for pandas DataFrame Series
    width  = int(row_data.iloc[10])
    height = int(row_data.iloc[11])
    subsampling = str(row_data.iloc[15])
    bitDepth = int(row_data.iloc[16]) if str(row_data.iloc[16]).strip().isdigit() else 8
    frames  = int(row_data.iloc[29])

    factor = _yuv_bytes_factor(subsampling)
    bytesize = int(width * height * factor * bitDepth / 8 * frames)

    # do some plausibility checks
    if datasize is None:
        print("sequence missing",
              {"row": row_index, "sequence": yuv_name, "expected_bytes": bytesize})
        datasize = bytesize
    elif datasize == bytesize:
        # 帧率/时长输出与原逻辑保持一致
        frameRate = float(row_data.iloc[17])
        duration  = float(row_data.iloc[3])
        print("size ok",
              {"row": row_index, "sequence": yuv_name, "bytes": datasize,
               "duration": duration, "frames/fps": frames / frameRate})
    else:
        # 输出一个“按实际大小反推估计帧数”的帮助值
        approx_frames = int(8 * datasize / (width * height * factor * bitDepth)) if datasize else 0
        print("size mismatch",
              {"row": row_index, "sequence": yuv_name,
               "actual": datasize, "expected": bytesize,
               "frameCount(csv)": frames, "approx_frames(from_size)": approx_frames})
    return datasize

# --- Main script execution ---
df = None

try:
    df = pd.read_csv('/content/drive/MyDrive/work/sequences.csv')
    print(f"Successfully read CSV from /content/drive/MyDrive/work/sequences.csv")
except Exception as e:
    print(f"Error reading CSV from /content/drive/MyDrive/work/sequences.csv: {e}")

if df is None:
    print("Error: Could not read the CSV file. Exiting.")
else:
    print(f"Total no. of rows: {len(df)}")
    output_dir = '/content/drive/MyDrive/work/Sequences' # Changed to Google Drive path
    os.makedirs(output_dir, exist_ok=True)

    for row_index, row_data in df.iterrows():
        json_filename = str(row_data.iloc[2]) + '.json' # Using 'Location' column for filename
        jsonFile = os.path.join(output_dir, json_filename)

        # Safely get specification_value, defaulting if index is out of range
        specification_value = ""
        try:
            # Access 'specification' column (index 42)
            # Use .get() with a default to avoid KeyError if column doesn't exist,
            # or check len before .iloc if sure about index.
            # Assuming 'specification' is always at index 42 as per user prompt.
            if len(row_data) > 42:
                specification_value = str(row_data.iloc[42])
        except IndexError:
            print(f"Warning: Row {row_index} is too short to contain specification_value at index 42. Using default empty string.")
        except KeyError: # Fallback if 'specification' is expected as named column but not found
            print(f"Warning: 'specification' column not found for row {row_index}. Using default empty string.")

        base_yuv_url = BASE_URL_CONFORMANCE # Default
        if "26955" in specification_value:
            base_yuv_url = BASE_URL_5GVIDEO
        elif "29265" in specification_value: # User mentioned 26955 or 29265
            base_yuv_url = BASE_URL_CONFORMANCE

        data = {
            "Sequence": {},
            "Signal": {
                "type": "",
                "Properties": {},
                "Components": []
            },
            "copyRight": "",
            "Contact": {}
        }

        data["Sequence"]["Name"] = str(row_data.iloc[0])

        reference_value = str(row_data.iloc[1])
        if "26955" in specification_value:
            data["Sequence"]["Specification"] = f"3GPP TR 26.955, {reference_value}"
        elif "29265" in specification_value:
            data["Sequence"]["Specification"] = f"3GPP TS 26.265, {reference_value}"
        else:
            data["Sequence"]["Specification"] = reference_value # Fallback if no specific match

        # Safely get 'Key' value, defaulting if index is out of range
        key_value = ""
        try:
            if len(row_data) > 41:
                key_value = str(row_data.iloc[41])
        except IndexError:
            print(f"Warning: Row {row_index} is too short to contain Key at index 41. Using default empty string.")
        data["Sequence"]["Key"] = key_value # Using index 41 for Key

        data["Sequence"]["Background"] = str(row_data.iloc[4])
        data["Signal"]["type"] = str(row_data.iloc[35]) # 'type' from index 35

        data["Signal"]["Properties"]["width"] = int(row_data.iloc[10])
        data["Signal"]["Properties"]["height"] = int(row_data.iloc[11])
        data["Signal"]["Properties"]["format"] = str(row_data.iloc[12])
        data["Signal"]["Properties"]["packing"] = str(row_data.iloc[13])
        data["Signal"]["Properties"]["scan"] = str(row_data.iloc[14])
        data["Signal"]["Properties"]["subsampling"] = str(row_data.iloc[15])
        data["Signal"]["Properties"]["bitDepth"] = int(row_data.iloc[16]) if str(row_data.iloc[16]).strip().isdigit() else 8
        data["Signal"]["Properties"]["frameRate"] = float(row_data.iloc[17])
        data["Signal"]["Properties"]["colourPrimaries"] = str(row_data.iloc[18])
        data["Signal"]["Properties"]["transferCharacteristics"] = str(row_data.iloc[19])
        data["Signal"]["Properties"]["matrixCoefficients"] = str(row_data.iloc[20])
        data["Signal"]["Properties"]["videoFullRangeFlag"] = str(row_data.iloc[21])
        data["Signal"]["Properties"]["chromaSampleLocType"] = str(row_data.iloc[22])
        data["Signal"]["Properties"]["sampleAspectRatio"] = str(row_data.iloc[23])
        data["Signal"]["Properties"]["duration"] = float(row_data.iloc[3])
        data["Signal"]["Properties"]["frameCount"] = int(row_data.iloc[29])
        data["Signal"]["Properties"]["startFrame"] = int(row_data.iloc[28])

        # Initialize variables for local file paths and HTTP URIs
        local_yuv_path = ""
        http_uri = ""
        datasize = 0
        md5_value = ""

        local_left_eye_path = ""
        http_left_eye_uri = ""
        leftEyeDatasize = 0
        leftEyemd5 = ""

        local_right_eye_path = ""
        http_right_eye_uri = ""
        rightEyeDatasize = 0
        rightEyemd5 = ""

        # 'Location' column (index 2) provides the subdirectory
        sequence_location_dir = str(row_data.iloc[2])

        # Mono or Stereo processing
        if data["Signal"]["type"] != "stereo":
            # For mono, construct paths using URI (index 6) and md5 (index 7)
            mono_filename = str(row_data.iloc[6])
            # Construct the full relative path including the subdirectory
            full_relative_path = os.path.join(sequence_location_dir, mono_filename)

            local_yuv_path = full_relative_path # Use this for local file checks
            # Corrected: Use string concatenation for URL instead of os.path.join
            http_uri = base_yuv_url + full_relative_path

            datasize = get_datasize(row_data, local_yuv_path, row_index)
            md5_from_csv = str(row_data.iloc[7])
            md5_value = get_md5_value(local_yuv_path, md5_from_csv)
        else:
            # For stereo, construct paths for left and right eyes
            # Left eye (first component) uses URI (index 6) and md5 (index 7)
            left_eye_filename = str(row_data.iloc[6])
            full_relative_left_eye_path = os.path.join(sequence_location_dir, left_eye_filename)

            local_left_eye_path = full_relative_left_eye_path # Use this for local file checks
            # Corrected: Use string concatenation for URL instead of os.path.join
            http_left_eye_uri = base_yuv_url + full_relative_left_eye_path

            leftEyeDatasize = get_datasize(row_data, local_left_eye_path, row_index)
            md5_from_csv_left = str(row_data.iloc[7])
            leftEyemd5 = get_md5_value(local_left_eye_path, md5_from_csv_left)

            # Right eye (second component) uses URI2 (index 43) and md52 (index 44)
            right_eye_filename = ""
            md5_from_csv_right = ""
            full_relative_right_eye_path = ""

            # Check if URI2 (index 43) and md52 (index 44) exist for the row
            if len(row_data) > 43: # Check if index 43 (URI2) exists
                right_eye_filename = str(row_data.iloc[43])
                full_relative_right_eye_path = os.path.join(sequence_location_dir, right_eye_filename)

                local_right_eye_path = full_relative_right_eye_path # Use this for local file checks
                # Corrected: Use string concatenation for URL instead of os.path.join
                http_right_eye_uri = base_yuv_url + full_relative_right_eye_path

                rightEyeDatasize = get_datasize(row_data, local_right_eye_path, row_index)

                if len(row_data) > 44: # Check if index 44 (md52) exists
                    md5_from_csv_right = str(row_data.iloc[44])
                rightEyemd5 = get_md5_value(local_right_eye_path, md5_from_csv_right)
            else:
                print(f"Warning: Row {row_index} is stereo but missing URI2 at index 43. Right eye component will be incomplete.")

        if data["Signal"]["type"] == "stereo":
            # Safely get 'flag_hero' value, defaulting if index is out of range
            flag_hero = ""
            try:
                if len(row_data) > 40:
                    flag_hero = str(row_data.iloc[40]) # Using index 40 for flag_hero
            except IndexError:
                print(f"Warning: Row {row_index} is too short to contain flag_hero at index 40. Using default empty string.")

            hero_left = "yes" if flag_hero == "left" else "no"
            hero_right = "yes" if flag_hero == "right" else "no"

            add_component(
                data,
                seq_type = data["Signal"]["type"],
                comp_type="left",
                hero=hero_left,
                uri=http_left_eye_uri,
                md5=leftEyemd5,
                size=leftEyeDatasize
            )
            # Always add right eye component for stereo, even if data is incomplete
            add_component(
                data,
                seq_type = data["Signal"]["type"],
                comp_type="right",
                hero=hero_right,
                uri=http_right_eye_uri,
                md5=rightEyemd5,
                size=rightEyeDatasize
            )
        else:
            add_component(
                data,
                seq_type = data["Signal"]["type"],
                comp_type="",
                hero="",
                uri=http_uri,
                md5=md5_value,
                size=datasize
            )

        data["copyRight"] = str(row_data.iloc[30])
        data["Contact"]["Name"] = str(row_data.iloc[31])
        data["Contact"]["Company"] = str(row_data.iloc[32])
        data["Contact"]["e-mail"] = str(row_data.iloc[33])
        data["Contact"]["generation"] = str(row_data.iloc[34])

        # delete first, if exist
        if os.path.exists(jsonFile):
            try:
                os.remove(jsonFile)
                print(f"[INFO] removed existing file: {jsonFile}")
            except Exception as e:
                print(f"[ERROR] failed to remove {jsonFile}: {e}")

        # dump json file
        with open(jsonFile, 'w') as outfile:
            json.dump(data, outfile, sort_keys=False, indent=4)