# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import shutil
import argparse
import shutil
import subprocess

from glob import glob

from anydata.utils.read import read_json
from anydata.utils.write import write_json
from anydata.utils.data import merge_dicts
from anydata.sync.sync_utils import aws_s3_cp, aws_s3_sync
from anydata.converters.utils import loop_over, crawl

#######################################################

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("path", type=str)
    parser.add_argument("--local_folder", type=str, default='/data/cv_unified')
    parser.add_argument("--s3_path", type=str, default='s3://tri-ml-sandbox-16011-us-west-2-datasets/cv_unified')
    parser.add_argument('--subfolder', type=str, default=None)
    parser.add_argument('--download', action='store_true')  # Download dataset for processing
    parser.add_argument('--upload', action='store_true')    # Upload split to s3
    parser.add_argument('--official', action='store_true')  # If it's not official, use "_debug"
    parser.add_argument('--execute', action='store_true')    # Execute changes (otherwise it's dryrun)
    parser.add_argument('--changes', type=str, nargs='+')   # Which changes will be made
    args = parser.parse_args()

    # Make it debug if not official
    if not args.official:
        args.local_folder += '_debug'
        args.s3_path += '_debug'

    # If path is not absolute, assume unified folder
    if args.path.startswith('/'):
        args.src = args.path
    else:
        args.src = f'{args.local_folder}/{args.path}'

    # Parse changes into a dictionary
    # Example: --changes depth/metric=False will set metadata['depth']['metric'] = False
    changes = {key: dict() for key in ['+', '-', '=']}
    for change in args.changes:
        if "=" in change:
            keys, val = change.split('=')
            mode = '=' 
        elif '+' in change:
            keys, val = change.split('+')
            mode = '+' 
        elif '-' in change:
            keys, val = change.split('-')
            mode = '-' 
        else:
            raise ValueError(f'Invalid change type {change}')
        keys = keys.split('/')
        rec_add(changes, keys, val, mode)
    args.changes = changes

    return args

#######################################################

def rec_add(data, keys, val, mode=None):
    """Recursive addition to metadata, going inside nested dictionaries"""
    if mode is not None:
        data = data[mode]
    if len(keys) > 1:
        if keys[0] not in data: data[keys[0]] = dict()
        rec_add(data[keys[0]], keys[1:], val)
    else:
        data[keys[0]] = val        

def coerce_value(val):
    '''
    Coerce a string change value to bool / int / float, else leave as string.
    '''
    if not isinstance(val, str):
        return val
    if val == 'True':
        return True
    if val == 'False':
        return False
    try:
        return int(val)
    except ValueError:
        pass
    try:
        return float(val)
    except ValueError:
        pass
    return val

def modify(args, metadata, changes, mode):
    """Modify metadata given the changes"""
    name = metadata.get('raw_id', metadata.get('info', {}).get('raw_id', 'SHARED'))
    for key1, val1 in changes.items():
        if isinstance(val1, dict):
            if key1 not in metadata: metadata[key1] = dict()
            for key2, val2 in val1.items():
                if isinstance(val2, dict):
                    if key2 not in metadata[key1]: metadata[key1][key2] = dict()
                    for key3, val3 in val2.items():
                        if key3 not in metadata[key1][key2]: metadata[key1][key2][key3] = None
                        print(f'### {name} : {key1} / {key2} / {key3}: {metadata[key1][key2][key3]} --> {val3}')
                        if mode == '=':
                            metadata[key1][key2][key3] = coerce_value(val3)
                        elif mode == '+' and val3 not in metadata[key1][key2][key3]: 
                            metadata[key1][key2][key3].append(val3)
                        elif mode == '-' and val3 in metadata[key1][key2][key3]: 
                            metadata[key1][key2][key3].remove(val3)
                else:
                    if key2 not in metadata[key1]: metadata[key1][key2] = None
                    print(f'### {name} : {key1} / {key2} : {metadata[key1][key2]} {mode} {val2}')
                    if mode == '=':
                        metadata[key1][key2] = coerce_value(val2)
                    elif mode == '+' and val2 not in metadata[key1][key2]: 
                        metadata[key1][key2].append(val2)
                    elif mode == '-' and val2 in metadata[key1][key2]: 
                        metadata[key1][key2].remove(val2)
        else:
            if key1 not in metadata: metadata[key1] = None
            print(f'### {name} : {key1} : {metadata[key1]} {mode} {val1}')
            if mode == '=':
                metadata[key1] = coerce_value(val1)
            elif mode == '+' and val1 not in metadata[key1]: 
                metadata[key1].append(val1)
            elif mode == '-' and val1 in metadata[key1]: 
                metadata[key1].remove(val1)

#######################################################

def change_metadata(args):

    ### Set paths
    s3_path = f'{args.s3_path}/{args.path}'
    local_path = f'{args.local_folder}/{args.path}'
    if args.subfolder is not None:
        s3_path += f'/{args.subfolder}'
        local_path += f'/{args.subfolder}'

    if args.download:
        ### Download dataset (METADATA only)
        aws_s3_sync(s3_path, local_path, 
            extras='--exclude "*" --include "*metadata.json" --include "*metadata_shared.json"'
        )

    # Find all relevant files
    sequences = read_json(f'{args.src}/split_all.json')['sequences']
    files = [f'{args.src}/{s}/metadata.json' for s in sequences]
    assert len(files) > 0, f'{args.src} not found!'

    print('########################################')
    print('##### CHANGING METADATA ...')
    print('########################################')
    for mode, changes in args.changes.items():
        for key1, val1 in changes.items():
            if isinstance(val1, dict):
                print(f'### {key1}')
                for key2, val2 in val1.items():
                    if isinstance(val2, dict):
                        print(f'###### {key2}')
                        for key3, val3 in val2.items():
                            print(f'######### {key3} : {val3}')
                    else:
                        print(f'###### {key2} : {val2}')
            else:
                print(f'### {key1} : {val1}')
    print('########################################')

    # Modify shared metadata
    file = f'{args.src}/metadata_shared.json'
    metadata = read_json(file)
    for mode, changes in args.changes.items():
        modify(args, metadata, changes, mode)
    if args.execute:
        write_json(file, metadata)

    # Modify sequence metadata
    for file in files:
        metadata = read_json(file)
        for mode, changes in args.changes.items():
            modify(args, metadata, changes, mode)
        if args.execute:
            write_json(file, metadata)

    if args.execute and args.upload:
        print(f'### UPLOADING {local_path} -> {s3_path}')
        aws_s3_sync(local_path, s3_path, 
            extras='--exclude "*" --include "*metadata.json" --include "*metadata_shared.json"'
        )

#######################################################

if __name__ == '__main__':
    args = parse_args()
    change_metadata(args)

#######################################################