#!/root/venv/bin/python
import json
import yaml
from collections import OrderedDict
from ruamel.yaml import YAML
import sys
from ruamel.yaml.scalarstring import PreservedScalarString


def detect_format(file_path):
    """Detect the file format based on the extension or content."""
    if file_path.endswith('.json'):
        return 'json'
    elif file_path.endswith(('.yaml', '.yml')):
        return 'yaml'
    else:
        # Try to detect the format by reading the file content
        with open(file_path, 'r') as file:
            content = file.read().strip()
            try:
                json.loads(content)
                return 'json'
            except json.JSONDecodeError:
                try:
                    yaml.safe_load(content)
                    return 'yaml'
                except yaml.YAMLError:
                    raise ValueError("Unable to detect format. Please use a valid JSON or YAML file.")


def load_data(file_path, format):
    """Load data from a JSON or YAML file."""
    with open(file_path, 'r') as file:
        if format == 'json':
            return json.load(file)
        elif format == 'yaml':
            return yaml.safe_load(file)
        else:
            raise ValueError("Unsupported format. Use JSON or YAML.")


def order_data(data, key_order):
    """Order data based on the provided key_order map."""
    if isinstance(data, dict):
        ordered = OrderedDict()
        for key, sub_order in key_order.items():
            if key in data:
                ordered[key] = order_data(data[key], sub_order)
        for key in data:
            if key not in key_order:
                ordered[key] = data[key]  # Include additional keys at the end
        return ordered
    elif isinstance(data, list):
        return [order_data(item, key_order[0]) for item in data] if key_order else data
    else:
        return data


def render_output(data, output_format):
    """Render the ordered data into JSON or YAML."""
    if output_format == "json":
        return json.dumps(data, indent=4)
    elif output_format == "yaml":
        yaml = YAML()
        yaml.default_flow_style = False
        yaml.indent(mapping=2, sequence=4, offset=2)
        
        # Use a lambda to correctly call represent_dict with the right arguments
        yaml.representer.add_representer(OrderedDict, lambda self, data: self.represent_dict(data))
        
        # Convert long strings to block style
        def convert_to_block_style(obj):
            if isinstance(obj, str) and '\n' in obj:
                return PreservedScalarString(obj)
            elif isinstance(obj, dict):
                return {k: convert_to_block_style(v) for k, v in obj.items()}
            elif isinstance(obj, list):
                return [convert_to_block_style(elem) for elem in obj]
            return obj
        
        data = convert_to_block_style(data)
        
        from io import StringIO
        stream = StringIO()
        yaml.dump(data, stream)
        return stream.getvalue()
    else:
        raise ValueError("Output format must be 'json' or 'yaml'.")


if __name__ == "__main__":
    # Define the key order map
    key_order = {
        "version": [],
        "sections": {
            "main": [
                {
                    "ai": {
                        "prompt": {
                            "confidence": [],
                            "barge_confidence": [],
                            "top_p": [],
                            "temperature": [],
                            "frequency_penalty": [],
                            "presence_penalty": [],
                            "text": [],
                            "contexts": {
                                "default": {
                                    "steps": [
                                        {
                                            "name": [],
                                            "text": [],
                                            "functions": [],
                                            "step_criteria": [],
                                            "valid_steps": [],
                                            "end": [],
                                            "skip_user_turn": []
                                        }
                                    ]
                                }
                            }
                        },
                        "params": {
                            "direction": [],
                            "wait_for_user": [],
                            "end_of_speech_timeout": [],
                            "attention_timeout": [],
                            "inactivity_timeout": [],
                            "outbound_attention_timeout": [],
                            "background_file": [],
                            "background_file_loops": [],
                            "background_file_volume": [],
                            "local_tz": [],
                            "conscience": [],
                            "ai_volume": [],
                            "save_conversation": [],
                            "conversation_id": [],
                            "digit_timeout": [],
                            "digit_terminators": [],
                            "energy_level": [],
                            "swaig_allow_swml": [],
                            "ai_model": [],
                            "audible_debug": []
                        },
                        "post_prompt_url": [],
                        "post_prompt": {
                            "confidence": [],
                            "barge_confidence": [],
                            "top_p": [],
                            "temperature": [],
                            "frequency_penalty": [],
                            "presence_penalty": [],
                            "text": []
                        },
                        "languages": [
                            {
                                "name": [],
                                "code": [],
                                "voice": []
                            }
                        ],
                        "hints": [],
                        "pronounce": [
                            {
                                "replace": [],
                                "with": [],
                                "ignore_case": []
                            }
                        ],
                        "SWAIG": {
                            "defaults": {
                                "web_hook_url": [],
                                "meta_data_token": [],
                                "meta_data": {
                                    "my_key": []
                                }
                            },
                            "native_functions": [],
                            "includes": [
                                {
                                    "url": [],
                                    "functions": []
                                }
                            ],
                            "functions": [
                                {
                                    "function": [],
                                    "meta_data_token": [],
                                    "meta_data": {
                                        "my_key": []
                                    },
                                    "purpose": [],
                                    "description": [],
                                    "argument": {
                                        "type": [],
                                        "properties": {
                                            "location": {
                                                "type": [],
                                                "description": []
                                            },
                                            "street": {
                                                "type": [],
                                                "description": []
                                            },
                                            "zipcode": {
                                                "type": [],
                                                "description": []
                                            },
                                            "genre": {
                                                "type": [],
                                                "description": []
                                            },
                                            "location_id": {
                                                "type": [],
                                                "description": []
                                            }
                                        }
                                    },
                                    "data_map": {
                                        "webhooks": [
                                            {
                                                "url": [],
                                                "headers": {
                                                    "X-RapidAPI-Key": [],
                                                    "X-RapidAPI-Host": []
                                                },
                                                "expressions": [
                                                    {
                                                        "string": [],
                                                        "pattern": [],
                                                        "output": {
                                                            "response": [],
                                                            "action": {
                                                                "set_meta_data": {
                                                                    "addr": [],
                                                                    "LL": []
                                                                }
                                                            }
                                                        }
                                                    }
                                                ],
                                                "output": {
                                                    "response": []
                                                }
                                            }
                                        ],
                                        "foreach": {
                                            "input_key": [],
                                            "output_key": [],
                                            "max": [],
                                            "append": []
                                        },
                                        "output": {
                                            "response": []
                                        }
                                    }
                                }
                            ]
                        }
                    }
                }
            ]
        }
    }

    # Input: specify file and output format
    input_file = sys.argv[1]  # JSON or YAML file path
    output_format = sys.argv[2].lower()  # "json" or "yaml"

    try:
        # Detect the format of the input file
        input_format = detect_format(input_file)

        # Load, order, and render
        data = load_data(input_file, input_format)
        ordered_data = order_data(data, key_order)
        output = render_output(ordered_data, output_format)
        print(output)
    except Exception as e:
        print(f"Error: {e}")