import yaml
import json
import re
import sys
import argparse
import os

# Regex to capture the type ID and anchor ID from the document separator
header_pattern = re.compile(r"--- !u!(\d+) &(\S+)")
# Regex to find and remove the tags for the parser
tag_remover_pattern = re.compile(r"!u!\d+\s")

def convert_unity_yaml_to_json(yaml_content):
    """
    Parses a Unity YAML file string, preserving fileID references, and returns a JSON string.
    """
    json_data = []
    
    # First, find all the original headers
    headers = header_pattern.findall(yaml_content)
    
    # Next, remove the problematic tags from the content so the parser doesn't fail
    sanitized_content = tag_remover_pattern.sub("", yaml_content)
    
    # Use the standard SafeLoader, as the tags are now gone
    documents = list(yaml.safe_load_all(sanitized_content))

    # The first document is the file info, which we can often skip if it's empty
    if documents and isinstance(documents[0], str) and 'YAML' in documents[0]:
        documents.pop(0)

    if len(headers) != len(documents):
        print(f"Warning: Mismatch between headers found ({len(headers)}) and documents parsed ({len(documents)}).", file=sys.stderr)

    for i, doc in enumerate(documents):
        if i < len(headers):
            type_id, anchor_id = headers[i]
            
            structured_doc = {
                'type_id': type_id,
                'anchor_id': anchor_id,
                'data': doc
            }
            json_data.append(structured_doc)
        else:
            # Append any extra docs without headers (should be rare in Unity files)
            json_data.append({'data': doc})

    # Use compact encoding for the final JSON
    return json.dumps(json_data)

def main():
    parser = argparse.ArgumentParser(description='Convert Unity YAML assets to JSON.')
    parser.add_argument('input_path', type=str, help='Absolute path to the input Unity asset file.')
    parser.add_argument('output_path', type=str, help='Absolute path for the output JSON file.')
    args = parser.parse_args()

    input_path = args.input_path
    output_path = args.output_path

    try:
        # Ensure the output directory exists
        output_dir = os.path.dirname(output_path)
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        with open(input_path, 'r') as f:
            content = f.read()

        json_output = convert_unity_yaml_to_json(content)

        with open(output_path, 'w') as f:
            f.write(json_output)

        print(f"Successfully converted '{input_path}' to '{output_path}'")

    except Exception as e:
        print(f"An error occurred: {e}", file=sys.stderr)
        sys.exit(1)

if __name__ == "__main__":
    main()