#!/usr/bin/env python3

import json
import base64
import yaml
import os
import re
import sys
import shutil
import logging
from typing import Dict, Any, Optional

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)

class ConfigError(Exception):
    """Custom exception for configuration-related errors"""
    pass

def validate_file_path(path: str, description: str) -> None:
    """Validate that a file path exists and is accessible."""
    if not path:
        raise ConfigError(f"{description} path is empty")
    if not os.path.exists(path):
        raise ConfigError(f"{description} path does not exist: {path}")
    if not os.access(path, os.R_OK):
        raise ConfigError(f"{description} path is not readable: {path}")

def validate_directory_path(path: str, description: str) -> None:
    """Validate that a directory path exists and is writable."""
    if not path:
        raise ConfigError(f"{description} path is empty")
    if not os.path.exists(path):
        try:
            os.makedirs(path, exist_ok=True)
        except OSError as e:
            raise ConfigError(f"Failed to create {description} directory: {e}")
    if not os.access(path, os.W_OK):
        raise ConfigError(f"{description} path is not writable: {path}")

def merge_yaml(base_yaml: Dict[str, Any], overlay_yaml: Dict[str, Any]) -> Dict[str, Any]:
    """
    Merge two YAML dictionaries, with overlay taking precedence.
    """
    if not isinstance(base_yaml, dict) or not isinstance(overlay_yaml, dict):
        raise ConfigError("Both base and overlay YAML must be dictionaries")
    
    result = base_yaml.copy()
    for key, value in overlay_yaml.items():
        if key in result and isinstance(result[key], dict) and isinstance(value, dict):
            result[key] = merge_yaml(result[key], value)
        else:
            result[key] = value
    return result

def split_credentials(credentials_path: str, output_dir: str) -> None:
    """
    Split credentials file into separate username and password files.
    Supports both formats:
    - username=value
    - username="value"
    """
    try:
        # Ensure output directory exists
        os.makedirs(output_dir, exist_ok=True)
        
        with open(credentials_path, 'r') as f:
            content = f.read()
        
        # Try to match both formats: username=value and username="value"
        username_match = re.search(r'username=(["\']?)([^"\'\s]+)\1', content)
        password_match = re.search(r'password=(["\']?)([^"\'\s]+)\1', content)
        
        if not username_match or not password_match:
            raise ConfigError("Invalid credentials format: missing username or password")
        
        username = username_match.group(2)
        password = password_match.group(2)
        
        if not username or not password:
            raise ConfigError("Empty username or password found in credentials")
        
        # Write username to file with new name
        username_path = os.path.join(output_dir, 'alertmanager-client-cred-username')
        with open(username_path, 'w') as f:
            f.write(username)
        os.chmod(username_path, 0o600)  # Set restrictive permissions
        
        # Write password to file with new name
        password_path = os.path.join(output_dir, 'alertmanager-client-cred-password')
        with open(password_path, 'w') as f:
            f.write(password)
        os.chmod(password_path, 0o600)  # Set restrictive permissions
        
    except IOError as e:
        raise ConfigError(f"Failed to process credentials: {e}")

def validate_input_data(input_data: Dict[str, Any]) -> None:
    """Validate the input JSON data structure."""
    required_fields = [
        'Overlay',
        'ClientCredentialsPath',
        'InitialConfigPath',
        'FinalClientCredentialsDir',
        'FinalConfigPath'
    ]
    
    for field in required_fields:
        if field not in input_data:
            raise ConfigError(f"Missing required field: {field}")
        if not input_data[field] and field != 'ClientCredentialsPath':
            raise ConfigError(f"Empty value for required field: {field}")

def main() -> None:
    try:
        if len(sys.argv) != 2:
            logging.error("Usage: python generate_prometheus_config_yaml.py <input_file>")
            sys.exit(1)

        input_file = sys.argv[1]
        validate_file_path(input_file, "Input JSON")
        
        # Read input JSON
        try:
            with open(input_file, 'r') as f:
                input_data = json.load(f)
        except json.JSONDecodeError as e:
            raise ConfigError(f"Invalid JSON in input file: {e}")
        
        # Validate input data structure
        validate_input_data(input_data)
        
        # Validate file paths
        if input_data['ClientCredentialsPath'] != "":
            validate_file_path(input_data['ClientCredentialsPath'], "Client credentials")
        validate_file_path(input_data['InitialConfigPath'], "Initial config")
        validate_directory_path(input_data['FinalClientCredentialsDir'], "Final client credentials")
        
        # Decode base64 overlay YAML
        try:
            overlay_yaml_str = base64.b64decode(input_data['Overlay']).decode('utf-8')
            overlay_yaml = yaml.safe_load(overlay_yaml_str)
            if not isinstance(overlay_yaml, dict):
                raise ConfigError("Overlay YAML must be a dictionary")
        except (base64.binascii.Error, UnicodeDecodeError) as e:
            raise ConfigError(f"Invalid base64 encoded overlay: {e}")
        except yaml.YAMLError as e:
            raise ConfigError(f"Invalid YAML in overlay: {e}")
        
        # Check if final config exists, if yes use it as base, otherwise use initial config
        base_config_path = input_data['FinalConfigPath'] if os.path.exists(input_data['FinalConfigPath']) else input_data['InitialConfigPath']
        validate_file_path(base_config_path, "Base config")
        
        # Read base config
        try:
            with open(base_config_path, 'r') as f:
                base_yaml = yaml.safe_load(f)
            if not isinstance(base_yaml, dict):
                raise ConfigError("Base YAML must be a dictionary")
        except yaml.YAMLError as e:
            raise ConfigError(f"Invalid YAML in base config: {e}")
        
        # Merge YAMLs
        final_yaml = merge_yaml(base_yaml, overlay_yaml)
        
        # Create backup of final config if it exists
        if os.path.exists(input_data['FinalConfigPath']):
            backup_file = input_data['FinalConfigPath'] + ".bk"
            try:
                shutil.copy2(input_data['FinalConfigPath'], backup_file)
                logging.info(f"Created backup of final config: {backup_file}")
            except OSError as e:
                raise ConfigError(f"Failed to create backup: {e}")
        
        # Write final config
        try:
            with open(input_data['FinalConfigPath'], 'w') as f:
                yaml.dump(final_yaml, f, default_flow_style=False)
            logging.info(f"Successfully wrote final config to: {input_data['FinalConfigPath']}")
        except IOError as e:
            raise ConfigError(f"Failed to write final config: {e}")
        
        # Split credentials
        if input_data['ClientCredentialsPath']  != "":
            try:
                split_credentials(input_data['ClientCredentialsPath'], input_data['FinalClientCredentialsDir'])
                logging.info(f"Successfully split credentials to: {input_data['FinalClientCredentialsDir']}")
            except ConfigError as e:
                raise ConfigError(f"Failed to split credentials: {e}")
        
        logging.info("Script completed successfully")
        
    except ConfigError as e:
        logging.error(f"Configuration error: {e}")
        sys.exit(1)
    except Exception as e:
        logging.error(f"Unexpected error: {e}")
        sys.exit(1)

if __name__ == '__main__':
    main()