#!/usr/bin/env python3
"""
Extract Internal Layer IDs and Update JSON Files

This script first extracts the internal layer IDs from a PSD file using ExtendScript,
then updates the corresponding JSON file with those IDs. This ensures that
when the API is used to update text layers, it has the correct internal IDs
that Adobe's API requires.

Usage:
    python extract_and_update_json.py [--psd-path PSD_PATH] [--json-path JSON_PATH]

If the PSD or JSON paths are not provided, the script will look for them.
"""

import os
import sys
import json
import argparse
import subprocess
import tempfile
from pathlib import Path
import glob
import time
import re

def create_jsx_script():
    """Create the ExtendScript to extract internal layer IDs"""
    return """
    // Extract Internal Layer IDs from Photoshop Document
    // This script extracts the actual internal IDs that Adobe's API uses
    // to identify text layers in a PSD file.
    #target photoshop

    // Function to extract text layers with their internal IDs
    function extractTextLayersWithInternalIDs() {
        if (!app.documents.length) {
            return {
                error: "No document open"
            };
        }
        
        var doc = app.activeDocument;
        var result = {
            documentName: doc.name,
            psdPath: doc.fullName.fsName,
            extractedAt: new Date().toString(),
            dimensions: {
                width: doc.width.as('px'),
                height: doc.height.as('px')
            },
            textLayerCount: 0,
            textLayers: []
        };
        
        // Utility function to traverse layers and extract info
        function traverseLayers(layers, path) {
            path = path || "";
            for (var i = 0; i < layers.length; i++) {
                var layer = layers[i];
                
                if (layer.typename === "LayerSet") {
                    // This is a layer group, traverse its children
                    var groupPath = path ? path + "/" + layer.name : layer.name;
                    traverseLayers(layer.layers, groupPath);
                } else if (layer.kind === LayerKind.TEXT) {
                    // This is a text layer, extract its information
                    extractTextLayerInfo(layer, path);
                }
            }
        }
        
        function extractTextLayerInfo(layer, path) {
            try {
                var layerPath = path ? path + "/" + layer.name : layer.name;
                
                // Get the layer's ID directly from the layer object
                var directID = layer.id;
                
                // Get the layer's internal ID using ActionManager
                var internalID = getInternalLayerID(layer);
                
                // Extract text content
                var text = layer.textItem.contents;
                
                // Extract basic style information
                var styleInfo = {
                    font: layer.textItem.font,
                    size: layer.textItem.size.value,
                    color: null, // Will be populated if available
                    alignment: getTextAlignment(layer)
                };
                
                // Try to extract text color
                try {
                    if (layer.textItem.color.rgb) {
                        styleInfo.color = [
                            layer.textItem.color.rgb.red,
                            layer.textItem.color.rgb.green,
                            layer.textItem.color.rgb.blue
                        ];
                    }
                } catch (e) {
                    // Color not defined or error accessing it
                }
                
                // Extract detailed formatting info if available
                var richTextInfo = extractRichTextInfo(layer);
                if (richTextInfo && richTextInfo.styles && richTextInfo.styles.length > 0) {
                    styleInfo.styles = richTextInfo.styles;
                    var hasRichTextFormatting = richTextInfo.styles.length > 1;
                } else {
                    styleInfo.styles = [];
                    var hasRichTextFormatting = false;
                }
                
                // Create layer info object with both direct and internal IDs
                var layerInfo = {
                    id: directID, // The ID we normally use
                    internalID: internalID, // The ID Adobe API might use
                    name: layer.name,
                    path: layerPath,
                    text: text,
                    visible: layer.visible,
                    styleInfo: styleInfo,
                    hasRichTextFormatting: hasRichTextFormatting
                };
                
                // Add to results
                result.textLayers.push(layerInfo);
                result.textLayerCount++;
                
            } catch (e) {
                // Log error but continue processing other layers
                $.writeln("Error processing layer '" + layer.name + "': " + e.message);
            }
        }
        
        // Function to get internal layer ID using ActionManager
        function getInternalLayerID(layer) {
            try {
                // Select the layer
                var idslct = charIDToTypeID("slct");
                var desc = new ActionDescriptor();
                var ref = new ActionReference();
                ref.putIdentifier(charIDToTypeID("Lyr "), layer.id);
                desc.putReference(charIDToTypeID("null"), ref);
                executeAction(idslct, desc, DialogModes.NO);
                
                // Get the layer's item index
                var ref = new ActionReference();
                ref.putProperty(charIDToTypeID("Prpr"), charIDToTypeID("ItmI"));
                ref.putEnumerated(charIDToTypeID("Lyr "), charIDToTypeID("Ordn"), charIDToTypeID("Trgt"));
                var itemIndexDesc = executeActionGet(ref);
                var itemIndex = itemIndexDesc.getInteger(charIDToTypeID("ItmI"));
                
                // Get the layer's ID
                var ref = new ActionReference();
                ref.putProperty(charIDToTypeID("Prpr"), stringIDToTypeID("layerID"));
                ref.putEnumerated(charIDToTypeID("Lyr "), charIDToTypeID("Ordn"), charIDToTypeID("Trgt"));
                var layerDesc = executeActionGet(ref);
                
                // Check if the layer has a layerID property
                if (layerDesc.hasKey(stringIDToTypeID("layerID"))) {
                    return layerDesc.getInteger(stringIDToTypeID("layerID"));
                } else {
                    // Fallback to another method of ID extraction
                    return layer.id; // Use direct ID as fallback
                }
            } catch (e) {
                $.writeln("Error getting internal ID for layer '" + layer.name + "': " + e.message);
                return layer.id; // Use direct ID if we can't get the internal ID
            }
        }
        
        // Function to get text alignment
        function getTextAlignment(layer) {
            try {
                var align = layer.textItem.justification;
                if (align === Justification.LEFT) {
                    return "left";
                } else if (align === Justification.CENTER) {
                    return "center";
                } else if (align === Justification.RIGHT) {
                    return "right";
                } else {
                    return "left"; // Default
                }
            } catch (e) {
                return "left"; // Default if error
            }
        }
        
        // Function to extract rich text information
        function extractRichTextInfo(layer) {
            try {
                var styles = [];
                
                // Get a reference to the layer
                var ref = new ActionReference();
                ref.putIdentifier(charIDToTypeID("Lyr "), layer.id);
                var desc = executeActionGet(ref);
                
                // Check if textKey exists
                if (desc.hasKey(stringIDToTypeID("textKey"))) {
                    var textKey = desc.getObjectValue(stringIDToTypeID("textKey"));
                    
                    // Check if textStyleRange exists
                    if (textKey.hasKey(stringIDToTypeID("textStyleRange"))) {
                        var styleRanges = textKey.getList(stringIDToTypeID("textStyleRange"));
                        
                        // Iterate through each style range
                        for (var i = 0; i < styleRanges.count; i++) {
                            var range = styleRanges.getObjectValue(i);
                            var from = range.getInteger(stringIDToTypeID("from"));
                            var to = range.getInteger(stringIDToTypeID("to"));
                            var styleRef = range.getObjectValue(stringIDToTypeID("textStyle"));
                            
                            var rangeText = layer.textItem.contents.substring(from, to);
                            var fontName = layer.textItem.font;
                            var fontSize = layer.textItem.size.value;
                            var fontStyle = "Regular";
                            var fontColor = null;
                            
                            // Extract font name
                            try {
                                if (styleRef.hasKey(stringIDToTypeID("fontName"))) {
                                    fontName = styleRef.getString(stringIDToTypeID("fontName"));
                                }
                            } catch (e) {}
                            
                            // Extract font style
                            try {
                                if (styleRef.hasKey(stringIDToTypeID("fontStyleName"))) {
                                    fontStyle = styleRef.getString(stringIDToTypeID("fontStyleName"));
                                }
                            } catch (e) {}
                            
                            // Extract font size
                            try {
                                if (styleRef.hasKey(stringIDToTypeID("size"))) {
                                    fontSize = styleRef.getUnitDoubleValue(stringIDToTypeID("size"));
                                }
                            } catch (e) {}
                            
                            // Extract color
                            try {
                                if (styleRef.hasKey(stringIDToTypeID("color"))) {
                                    var colorObj = styleRef.getObjectValue(stringIDToTypeID("color"));
                                    var colorValues = colorObj.getObjectValue(stringIDToTypeID("color"));
                                    
                                    // Check color model
                                    if (colorValues.hasKey(stringIDToTypeID("red"))) {
                                        fontColor = [
                                            Math.round(colorValues.getDouble(stringIDToTypeID("red"))),
                                            Math.round(colorValues.getDouble(stringIDToTypeID("green"))),
                                            Math.round(colorValues.getDouble(stringIDToTypeID("blue")))
                                        ];
                                    }
                                }
                            } catch (e) {}
                            
                            // Add style info to the array
                            styles.push({
                                start: from,
                                end: to,
                                text: rangeText,
                                font: fontName,
                                style: fontStyle,
                                size: fontSize,
                                color: fontColor
                            });
                        }
                    }
                }
                
                return { styles: styles };
            } catch (e) {
                $.writeln("Error extracting rich text info: " + e.message);
                return { styles: [] };
            }
        }
        
        // Start the layer traversal
        traverseLayers(doc.layers);
        
        return result;
    }

    // Main function to extract IDs and write to a file
    function main() {
        try {
            var result = extractTextLayersWithInternalIDs();
            
            // Convert result to formatted JSON
            var jsonString = JSON.stringify(result, null, 2);
            
            // Create a temp file to save the result
            var tempFolder = Folder.temp;
            var fileName = "ps_layer_ids_" + new Date().getTime() + ".json";
            var outputFile = new File(tempFolder + "/" + fileName);
            
            if (outputFile.open("w")) {
                outputFile.write(jsonString);
                outputFile.close();
                alert("Layer ID information saved to:\\n" + outputFile.fsName);
                return outputFile.fsName; // Return the path so the script can find it
            } else {
                alert("Error saving output file");
                return "";
            }
        } catch (e) {
            alert("Error: " + e.message);
            return "";
        }
    }

    // Run the script and return the path to the temp file
    main();
    """

def run_photoshop_script(psd_path):
    """
    Run the ExtendScript in Photoshop to extract layer IDs
    
    Args:
        psd_path: Path to the PSD file
        
    Returns:
        Path to the temporary JSON file with layer IDs
    """
    print(f"Opening PSD file in Photoshop: {os.path.basename(psd_path)}")
    
    # Create a temporary JSX file
    jsx_content = create_jsx_script()
    
    with tempfile.NamedTemporaryFile(suffix='.jsx', delete=False, mode='w') as jsx_file:
        jsx_file.write(jsx_content)
        jsx_path = jsx_file.name
    
    try:
        # Create AppleScript to run the JSX in Photoshop
        applescript = f"""
        tell application "Adobe Photoshop"
            activate
            open (POSIX file "{psd_path}")
            do javascript (POSIX file "{jsx_path}")
            close current document saving no
        end tell
        """
        
        # Save the AppleScript to a temporary file
        with tempfile.NamedTemporaryFile(suffix='.scpt', delete=False, mode='w') as as_file:
            as_file.write(applescript)
            as_path = as_file.name
        
        # Run the AppleScript
        print("Running ExtendScript to extract internal layer IDs...")
        result = subprocess.run(
            ["osascript", as_path],
            capture_output=True,
            text=True
        )
        
        if result.returncode != 0:
            print(f"Error running AppleScript: {result.stderr}")
            return None
        
        # Parse the output to find the path to the temporary JSON file
        output = result.stdout.strip()
        
        # Extract the path from the alert message
        match = re.search(r'Layer ID information saved to:\\n(.*)', output)
        if match:
            temp_json_path = match.group(1)
            print(f"Extracted layer IDs saved to temp file: {temp_json_path}")
            return temp_json_path
        else:
            print("Could not find the path to the temporary JSON file in the output")
            print(f"Output was: {output}")
            return None
        
    except Exception as e:
        print(f"Error running Photoshop script: {str(e)}")
        return None
    finally:
        # Clean up the temporary JSX file
        if os.path.exists(jsx_path):
            os.remove(jsx_path)
        
        # Clean up the temporary AppleScript file
        if 'as_path' in locals() and os.path.exists(as_path):
            os.remove(as_path)

def find_matching_json(psd_path):
    """
    Find the corresponding JSON file for a PSD file
    
    Args:
        psd_path: Path to the PSD file
        
    Returns:
        Path to the matching JSON file, or None if not found
    """
    psd_dir = os.path.dirname(psd_path)
    psd_name = os.path.basename(psd_path).replace('.psd', '')
    
    # Look for JSON files that match the PSD name
    json_patterns = [
        f"{psd_name}-textonly.json",
        f"{psd_name}-textonly-updated.json"
    ]
    
    for pattern in json_patterns:
        json_path = os.path.join(psd_dir, pattern)
        if os.path.exists(json_path):
            return json_path
    
    # If no direct match, search with glob
    json_files = glob.glob(os.path.join(psd_dir, f"{psd_name}*-textonly*.json"))
    
    if json_files:
        return json_files[0]
    
    return None

def find_matching_psd(json_path):
    """
    Find the corresponding PSD file for a JSON file
    
    Args:
        json_path: Path to the JSON file
        
    Returns:
        Path to the matching PSD file, or None if not found
    """
    json_dir = os.path.dirname(json_path)
    json_name = os.path.basename(json_path)
    
    # Strip common suffixes to get PSD name
    psd_name = json_name.replace('-textonly-updated.json', '').replace('-textonly.json', '')
    
    # Check if PSD exists directly
    psd_path = os.path.join(json_dir, f"{psd_name}.psd")
    if os.path.exists(psd_path):
        return psd_path
    
    # Read JSON file to get document name
    try:
        with open(json_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
            
        if 'documentName' in data:
            psd_path = os.path.join(json_dir, data['documentName'])
            if os.path.exists(psd_path):
                return psd_path
        
        if 'psdPath' in data:
            psd_path = data['psdPath']
            if os.path.exists(psd_path):
                return psd_path
            
            # Try to make path absolute
            if psd_path.startswith('~'):
                psd_path = os.path.expanduser(psd_path)
                if os.path.exists(psd_path):
                    return psd_path
                    
    except Exception as e:
        print(f"Error reading JSON file {json_path}: {str(e)}")
    
    # If still not found, try globbing
    psd_files = glob.glob(os.path.join(json_dir, f"{psd_name}*.psd"))
    
    if psd_files:
        return psd_files[0]
    
    return None

def update_json_with_internal_ids(json_path, layer_ids_data):
    """
    Update JSON file with internal layer IDs
    
    Args:
        json_path: Path to the JSON file to update
        layer_ids_data: Data from ExtendScript with internal layer IDs
        
    Returns:
        Path to the updated JSON file
    """
    print(f"Updating JSON file with internal layer IDs: {os.path.basename(json_path)}")
    
    try:
        # Load the existing JSON file
        with open(json_path, 'r', encoding='utf-8') as f:
            json_data = json.load(f)
        
        # Create a dictionary of layers from the layer_ids_data for easy lookup
        internal_layer_dict = {}
        
        for layer in layer_ids_data.get('textLayers', []):
            name = layer.get('name', '')
            internal_id = layer.get('internalID')
            direct_id = layer.get('id')
            
            if name and internal_id:
                internal_layer_dict[name] = {
                    'internalID': internal_id,
                    'directID': direct_id
                }
        
        # Update each layer in the JSON data
        for layer in json_data.get('textLayers', []):
            name = layer.get('name', '')
            
            if name in internal_layer_dict:
                # Add the internal ID
                layer['internalID'] = internal_layer_dict[name]['internalID']
                
                # Keep the original ID but mark it as directID if not already present
                if 'directID' not in layer:
                    layer['directID'] = layer.get('id')
                
                print(f"Layer '{name}': Added internal ID {layer['internalID']}")
        
        # Save to a new file with suffix -api-ready.json
        base_name = os.path.basename(json_path)
        dir_name = os.path.dirname(json_path)
        
        # Remove existing suffixes and add the new one
        new_name = base_name.replace('-textonly-updated.json', '').replace('-textonly.json', '')
        new_name = f"{new_name}-api-ready.json"
        
        new_path = os.path.join(dir_name, new_name)
        
        with open(new_path, 'w', encoding='utf-8') as f:
            json.dump(json_data, f, indent=2)
        
        print(f"Updated JSON file saved as: {new_path}")
        return new_path
        
    except Exception as e:
        print(f"Error updating JSON file: {str(e)}")
        return None

def process_file_pair(psd_path, json_path):
    """
    Process a PSD and JSON file pair to extract and update layer IDs
    
    Args:
        psd_path: Path to the PSD file
        json_path: Path to the JSON file
        
    Returns:
        Path to the updated JSON file
    """
    print("\n" + "="*70)
    print(f"Processing PSD: {os.path.basename(psd_path)}")
    print(f"With JSON: {os.path.basename(json_path)}")
    print("="*70 + "\n")
    
    # Extract layer IDs from the PSD file
    temp_json_path = run_photoshop_script(psd_path)
    
    if not temp_json_path:
        print("Failed to extract layer IDs from PSD file")
        return None
    
    # Load the layer ID data
    try:
        with open(temp_json_path, 'r', encoding='utf-8') as f:
            layer_ids_data = json.load(f)
    except Exception as e:
        print(f"Error loading layer ID data: {str(e)}")
        return None
    
    # Update the JSON file with the internal IDs
    updated_json_path = update_json_with_internal_ids(json_path, layer_ids_data)
    
    # Clean up the temporary file
    try:
        if os.path.exists(temp_json_path):
            os.remove(temp_json_path)
    except:
        pass
    
    return updated_json_path

def find_all_json_files(directory):
    """
    Find all JSON files with -textonly or -textonly-updated suffixes in a directory
    
    Args:
        directory: The directory to search
        
    Returns:
        List of JSON file paths
    """
    json_files = []
    patterns = ["*-textonly.json", "*-textonly-updated.json"]
    
    for pattern in patterns:
        json_files.extend(glob.glob(os.path.join(directory, pattern)))
    
    return json_files

def process_directory(directory):
    """
    Process all JSON files in a directory
    
    Args:
        directory: The directory to process
        
    Returns:
        List of updated JSON file paths
    """
    print(f"Scanning directory for JSON files: {directory}")
    
    json_files = find_all_json_files(directory)
    
    if not json_files:
        print("No JSON files found in the directory")
        return []
    
    print(f"Found {len(json_files)} JSON files to process")
    
    updated_files = []
    
    for json_path in json_files:
        # Find the corresponding PSD file
        psd_path = find_matching_psd(json_path)
        
        if psd_path:
            # Process the file pair
            updated_json = process_file_pair(psd_path, json_path)
            
            if updated_json:
                updated_files.append(updated_json)
        else:
            print(f"Could not find matching PSD file for {os.path.basename(json_path)}")
    
    return updated_files

def main():
    """Main function"""
    parser = argparse.ArgumentParser(
        description="Extract internal layer IDs and update JSON files for Adobe API",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Process a specific PSD and JSON file:
  python extract_and_update_json.py --psd-path /path/to/file.psd --json-path /path/to/file-textonly.json

  # Process all JSON files in a directory:
  python extract_and_update_json.py --directory /path/to/files

  # Process files in the current directory:
  python extract_and_update_json.py
        """
    )
    
    parser.add_argument('--psd-path', help='Path to the PSD file')
    parser.add_argument('--json-path', help='Path to the JSON file')
    parser.add_argument('--directory', help='Directory containing JSON files to process')
    
    args = parser.parse_args()
    
    # Check arguments and determine what to process
    if args.psd_path and args.json_path:
        # Process a specific PSD and JSON file
        if not os.path.exists(args.psd_path):
            print(f"Error: PSD file not found: {args.psd_path}")
            return
        
        if not os.path.exists(args.json_path):
            print(f"Error: JSON file not found: {args.json_path}")
            return
        
        updated_json = process_file_pair(args.psd_path, args.json_path)
        
        if updated_json:
            print(f"\nSuccessfully updated JSON file with internal layer IDs: {updated_json}")
            print("Use this file for API requests to ensure correct layer identification.")
        
    elif args.directory:
        # Process all JSON files in a directory
        if not os.path.isdir(args.directory):
            print(f"Error: Directory not found: {args.directory}")
            return
        
        updated_files = process_directory(args.directory)
        
        if updated_files:
            print(f"\nSuccessfully updated {len(updated_files)} JSON files with internal layer IDs:")
            for file_path in updated_files:
                print(f"  - {os.path.basename(file_path)}")
            print("\nUse these files for API requests to ensure correct layer identification.")
        else:
            print("\nNo JSON files were updated. Check the directory path and file names.")
        
    elif args.psd_path:
        # Process a specific PSD file but find the matching JSON
        if not os.path.exists(args.psd_path):
            print(f"Error: PSD file not found: {args.psd_path}")
            return
        
        json_path = find_matching_json(args.psd_path)
        
        if json_path:
            updated_json = process_file_pair(args.psd_path, json_path)
            
            if updated_json:
                print(f"\nSuccessfully updated JSON file with internal layer IDs: {updated_json}")
                print("Use this file for API requests to ensure correct layer identification.")
        else:
            print(f"Could not find matching JSON file for {os.path.basename(args.psd_path)}")
        
    elif args.json_path:
        # Process a specific JSON file but find the matching PSD
        if not os.path.exists(args.json_path):
            print(f"Error: JSON file not found: {args.json_path}")
            return
        
        psd_path = find_matching_psd(args.json_path)
        
        if psd_path:
            updated_json = process_file_pair(psd_path, args.json_path)
            
            if updated_json:
                print(f"\nSuccessfully updated JSON file with internal layer IDs: {updated_json}")
                print("Use this file for API requests to ensure correct layer identification.")
        else:
            print(f"Could not find matching PSD file for {os.path.basename(args.json_path)}")
        
    else:
        # Process the current directory
        updated_files = process_directory(os.getcwd())
        
        if updated_files:
            print(f"\nSuccessfully updated {len(updated_files)} JSON files with internal layer IDs:")
            for file_path in updated_files:
                print(f"  - {os.path.basename(file_path)}")
            print("\nUse these files for API requests to ensure correct layer identification.")
        else:
            print("\nNo JSON files were updated. Check the directory path and file names.")

if __name__ == "__main__":
    main()