#!/usr/bin/env python3

import os
import sys
import re
from merge_log4j2_yaml import merge_log4j2_yaml

component_name = sys.argv[1]
pod_id = sys.argv[2]

def merge_log4j2_with_overrides(base_file, overrides_file):
    """
    Merges base Log4j2 YAML file with overrides if they exist.

    Args:
        base_file: Path to the base Log4j2 YAML file
        overrides_file: Path to the overlay/overrides Log4j2 YAML file
        output_dir: Directory to write temporary merged file

    Returns:
        str: Path to the final output file (either original base_file if no merge
             occurred, or path to the successfully merged file)
    """
    # If overrides file doesn't exist, just return the base file
    if not os.path.exists(overrides_file):
        print(f"No Log4j2 overrides file found at {overrides_file}")
        return base_file

    print(f"Found Log4j2 overrides file: {overrides_file}")

    output_dir = os.environ.get('LOG4J2_DIR')
    # Create a temporary file for the merged result
    merged_output = os.path.join(output_dir, "merged_log4j2.yaml")

    # Try to merge the files
    try:
        success = merge_log4j2_yaml(base_file, overrides_file, merged_output)

        if success:
            print(f"Successfully merged base config with overrides")
            return merged_output
        else:
            print(f"WARNING: Failed to merge with overrides, using base config only")
            # Clean up the failed merge output if it exists
            if os.path.exists(merged_output):
                try:
                    os.remove(merged_output)
                except Exception:
                    pass
            return base_file

    except Exception as e:
        print(f"ERROR: Exception during YAML merge: {e}", file=sys.stderr)
        print(f"WARNING: Using base config only")
        if os.path.exists(merged_output):
            try:
                os.remove(merged_output)
            except Exception:
                pass
        return base_file

def process_log4j2_config():
    """
    Processes log4j2.yaml template using specific environment variables for substitution.
    Reads CP_COMPONENT_NAME env var to find inputs under /mnt/config.
    Writes the final config to /mnt/config/etc/<component>/log4j2.yaml.
    Substitutes ${app}, ${clusterId}, ${server} using APP_NAME, CLUSTER_ID, SERVER_ID env vars.
    """
    print("--- log4j2_config.py starting (using specific env vars: APP_NAME, CLUSTER_ID, SERVER_ID) ---")

    # Get necessary environment variables
    cp_component_name = os.environ.get('CP_COMPONENT_NAME')
    base_config_dir = "/mnt/config" # Standard base config mount path

    if not cp_component_name:
        print("ERROR: CP_COMPONENT_NAME environment variable not set.", file=sys.stderr)
        sys.exit(1)

    print(f"Running for component: {cp_component_name}")

    # Define INPUT paths based on component name and shared config
    component_config_dir = os.path.join(base_config_dir, cp_component_name)
    shared_config_dir = os.path.join(base_config_dir, "shared")

    log4j2_tmpl_source_file = os.path.join(component_config_dir, "log4j2.yaml")
    log4j2_shared_tmpl_source_file = os.path.join(shared_config_dir, "log4j2.yaml")
    
    # Define OUTPUT path within the shared /mnt/config volume
    log4j2_output_dir = os.environ.get('LOG4J2_DIR')
    log4j2_output_file = os.path.join(log4j2_output_dir, "log4j2.yaml") 

    # Determine template source (component-specific first, then shared)
    actual_tmpl_source = ""
    print(f"DEBUG: Checking for template file at '{log4j2_tmpl_source_file}' (exists: {os.path.exists(log4j2_tmpl_source_file)})") # DEBUG
    if os.path.exists(log4j2_tmpl_source_file):
        actual_tmpl_source = log4j2_tmpl_source_file
        print(f"Using component template source: {actual_tmpl_source}")
    else:
        print(f"DEBUG: Checking for template file at '{log4j2_shared_tmpl_source_file}' (exists: {os.path.exists(log4j2_shared_tmpl_source_file)})") # DEBUG
        if os.path.exists(log4j2_shared_tmpl_source_file):
            actual_tmpl_source = log4j2_shared_tmpl_source_file
            print(f"Using shared template source: {actual_tmpl_source}")
        else:
            print(f"WARNING: log4j2.yaml template source not found in {log4j2_tmpl_source_file} or {log4j2_shared_tmpl_source_file}. Exiting gracefully.")
            sys.exit(0)

    # Read template content
    template_content = "" # Initialize
    try:
        print(f"DEBUG: Attempting to read template file '{actual_tmpl_source}'")
        with open(actual_tmpl_source, 'r') as f_tmpl:
            template_content = f_tmpl.read()
        print(f"DEBUG: Successfully read template file. Length: {len(template_content)}")
    except Exception as e:
        print(f"ERROR: Failed to read template file {actual_tmpl_source}: {e}", file=sys.stderr)
        sys.exit(1)

    # --- Substitute properties using Specific Environment Variables --- 
    print("Substituting properties using specific environment variables...")

    namespace = os.environ['POD_NAMESPACE']

    replacements = {
        '${app}': component_name,
        '${clusterId}': namespace,
        '${server}': pod_id
    }

    processed_count = 0
    for placeholder, value in replacements.items():
        old_len = len(template_content)
        template_content = template_content.replace(placeholder, value)
        if len(template_content) != old_len:
            print(f"DEBUG: Replacement occurred for placeholder '{placeholder}'")
            processed_count += 1
        # else:
            # print(f"DEBUG: Placeholder '{placeholder}' not found in template content.")

    print(f"DEBUG: Finished processing specific environment variables. Performed {processed_count} replacements.") # DEBUG
    # --- End Environment Variable Substitution ---

    # Write the output file
    try:
        print(f"DEBUG: Attempting to write final config to '{log4j2_output_file}'") # DEBUG
        # Ensure the target directory exists
        print(f"DEBUG: Ensuring directory '{log4j2_output_dir}' exists.") # DEBUG
        os.makedirs(log4j2_output_dir, exist_ok=True)
        with open(log4j2_output_file, 'w') as f_out:
            f_out.write(template_content)
        print(f"Log4j2 configuration written to {log4j2_output_file}")
        print(f"DEBUG: File '{log4j2_output_file}' exists after write: {os.path.exists(log4j2_output_file)}") # DEBUG
    except Exception as e:
        print(f"ERROR: Failed to write final Log4j2 config to {log4j2_output_file}: {e}", file=sys.stderr)
        sys.exit(1)

    # Define path for overlay config
    log4j2_overrides_file = os.path.join(shared_config_dir, "log4j2_overrides.yaml")
    # Merge with overrides file if it exists
    final_output_file = merge_log4j2_with_overrides(log4j2_output_file, log4j2_overrides_file)
    # Move the final output to the expected location if it's different
    if final_output_file != log4j2_output_file:
        try:
            os.replace(final_output_file, log4j2_output_file)
            print(f"Final merged configuration written to {log4j2_output_file}")
        except Exception as e:
            print(f"ERROR: Failed to move merged config to final location: {e}", file=sys.stderr)
            return log4j2_output_file

    return log4j2_output_file

if __name__ == "__main__":
    process_log4j2_config()
    print("--- log4j2_config.py finished ---")
