import argparse
import json
import os.path
import subprocess

# In general where the resource type has changed, the new resource name is the same as the old resource.
# However in a few cases we have had to change the resource name to avoid conflicts with existing resources.
REMAPPED_RESOURCE_NAMES = {
    'db_create_schema_role_ownership': 'db_create_schema_account_role_ownership',
    'db_read_role_ownership': 'db_read_account_role_ownership',
    'db_read_role_all_view_in_schema_grants': 'db_read_role_all_non_materialized_view_in_schema_grants',
    'db_read_role_future_view_in_schema_grants': 'db_read_role_future_non_materialized_view_in_schema_grants',
    'db_read_role_all_view_in_database_grants': 'db_read_role_all_non_materialized_view_in_database_grants',
    'db_read_role_future_view_in_database_grants': 'db_read_role_future_non_materialized_view_in_database_grants',
    'db_readwrite_role_ownership': 'db_readwrite_account_role_ownership',
    'schema_read_role_ownership': 'schema_read_account_role_ownership',
    'schema_readwrite_role_ownership': 'schema_readwrite_account_role_ownership',
}


def main():
    project_path = get_project_path()
    terraform_state = get_terraform_state(project_path)

    imports = []
    removals = set()

    for resource in terraform_state['resources']:
        module = resource['module']
        resource_name = resource['name']
        module_without_index_key = module.split('[')[0]

        if resource['type'] == 'snowflake_role_ownership_grant':
            removals.add(f'{module_without_index_key}.snowflake_role_ownership_grant.{resource_name}')
            imports.extend(get_snowflake_role_ownership_grant_imports(module, resource))

        if resource['type'] == 'snowflake_warehouse_grant':
            removals.add(f'{module_without_index_key}.snowflake_warehouse_grant.{resource_name}')
            imports.extend(get_snowflake_warehouse_grant_imports(module, resource))

        if resource['type'] == 'snowflake_role_grants':
            removals.add(f'{module_without_index_key}.snowflake_role_grants.{resource_name}')
            imports.extend(get_snowflake_role_grants_imports(module, resource))

        if resource['type'] == 'snowflake_grant_privileges_to_role':
            removals.add(f'{module_without_index_key}.snowflake_grant_privileges_to_role.{resource_name}')
            imports.extend(get_snowflake_grant_privileges_to_role_imports(module, resource))

    write_imports_file(imports, project_path)
    write_removed_file(project_path, removals)


def get_project_path():
    parser = argparse.ArgumentParser(description='Tool for upgrading to terraform-snowflake v5.')
    parser.add_argument("path", help="Path to the directory to upgrade")
    args = parser.parse_args()
    path = args.path

    return os.path.abspath(path)


def get_terraform_state(project_path):
    try:
        result = subprocess.run(
            ["terraform", "state", "pull"],
            check=True,
            cwd=project_path,
            capture_output=True
        )
    except subprocess.CalledProcessError as e:
        raise Exception(
            f'''terraform state pull failed with exit code {e.returncode}. 
        Stdout: {e.stdout.decode("utf-8")}. 
        Stderr: {e.stderr.decode("utf-8")}''')
    terraform_state = json.loads(result.stdout.decode('utf-8'))
    return terraform_state


def get_new_resource_address(module, new_resource_type, name, index_key=None):
    remapped_name = REMAPPED_RESOURCE_NAMES.get(name, name)
    base_address = f'{module}.{new_resource_type}.{remapped_name}'

    if index_key is not None:
        if isinstance(index_key, str):
            return f'{base_address}["{index_key}"]'
        else:
            return f'{base_address}[{index_key}]'
    else:
        return base_address


def get_snowflake_role_ownership_grant_imports(module, resource):
    imports = []

    for instance in resource['instances']:
        on_role_name = instance['attributes']['on_role_name']
        to_role_name = instance['attributes']['to_role_name']
        to_resource = get_new_resource_address(
            module, 'snowflake_grant_ownership', resource['name'], instance.get('index_key'))
        imports.append({
            'to': to_resource,
            'id': f'ToAccountRole|{to_role_name}|COPY|OnObject|ROLE|{on_role_name}'
        })

    return imports


def get_snowflake_warehouse_grant_imports(module, resource):
    role_name = resource['instances'][0]['attributes']['roles'][0]
    warehouse_name = resource['instances'][0]['attributes']['warehouse_name']
    to_resource = get_new_resource_address(
        module, 'snowflake_grant_privileges_to_account_role', resource['name'])

    return [{
        'to': to_resource,
        'id': f'{role_name}|false|false|MONITOR,OPERATE,USAGE|OnAccountObject|WAREHOUSE|{warehouse_name}'
    }]


def get_snowflake_role_grants_imports(module, resource):
    imports = []

    for instance in resource['instances']:
        on_role_name = instance['attributes']['role_name']
        to_role_name = instance['attributes']['roles'][0]
        to_resource = get_new_resource_address(
            module, 'snowflake_grant_account_role', resource['name'], instance.get('index_key'))
        imports.append({
            'to': to_resource,
            'id': f'"{on_role_name}"|ROLE|"{to_role_name}"'
        })

    return imports


def get_snowflake_grant_privileges_to_role_imports(module, resource):
    imports = []

    for instance in resource['instances']:
        role_name = instance['attributes']['role_name']
        privileges = ','.join(instance['attributes']['privileges'])
        to_resource = get_new_resource_address(
            module, 'snowflake_grant_privileges_to_account_role', resource['name'], instance.get('index_key'))
        resource_id = None

        if len(instance['attributes']['on_account_object']) > 0:
            object_type = instance['attributes']['on_account_object'][0]['object_type']
            object_name = instance['attributes']['on_account_object'][0]['object_name']
            resource_id = f'{role_name}|false|false|{privileges}|OnAccountObject|{object_type}|{object_name}'

        elif len(instance['attributes']['on_schema']) > 0:
            resource_id = get_on_schema_grant_resource_id(instance, privileges, role_name)

        elif len(instance['attributes']['on_schema_object']) > 0:
            resource_id = get_on_schema_object_resource_id(instance, privileges, role_name)

        if not resource_id:
            raise Exception(f'Could not determine resource_id for {module}.snowflake_grant_privileges_to_account_role.{resource["name"]}')

        imports.append({
            'to': to_resource,
            'id': resource_id
        })

    return imports


def get_on_schema_grant_resource_id(instance, privileges, role_name):
    schema_name = instance['attributes']['on_schema'][0]['schema_name']
    all_schemas_in_database = instance['attributes']['on_schema'][0]['all_schemas_in_database']
    future_schemas_in_database = instance['attributes']['on_schema'][0]['future_schemas_in_database']
    if schema_name:
        return f'{role_name}|false|false|{privileges}|OnSchema|OnSchema|{schema_name}'
    elif all_schemas_in_database:
        return f'{role_name}|false|false|{privileges}|OnSchema|OnAllSchemasInDatabase|{all_schemas_in_database}'
    elif future_schemas_in_database:
        return f'{role_name}|false|false|{privileges}|OnSchema|OnFutureSchemasInDatabase|{future_schemas_in_database}'


def get_on_schema_object_resource_id(instance, privileges, role_name):
    object_type = instance['attributes']['on_schema_object'][0]['object_type']
    object_name = instance['attributes']['on_schema_object'][0]['object_name']
    all = instance['attributes']['on_schema_object'][0]['all']
    future = instance['attributes']['on_schema_object'][0]['future']
    if object_name:
        return f'{role_name}|false|false|{privileges}|OnSchemaObject|OnObject|{object_type}|{object_name}'
    elif all:
        in_database = all[0]['in_database']
        in_schema = all[0]['in_schema']
        object_type_plural = all[0]['object_type_plural']

        if in_database:
            return f'{role_name}|false|false|{privileges}|OnSchemaObject|OnAll|{object_type_plural}|InDatabase|{in_database}'
        elif in_schema:
            return f'{role_name}|false|false|{privileges}|OnSchemaObject|OnAll|{object_type_plural}|InSchema|{in_schema}'
    elif future:
        in_database = future[0]['in_database']
        in_schema = future[0]['in_schema']
        object_type_plural = future[0]['object_type_plural']

        if in_database:
            return f'{role_name}|false|false|{privileges}|OnSchemaObject|OnFuture|{object_type_plural}|InDatabase|{in_database}'
        elif in_schema:
            return f'{role_name}|false|false|{privileges}|OnSchemaObject|OnFuture|{object_type_plural}|InSchema|{in_schema}'

    return None


def write_imports_file(imports, project_path):
    imports.sort(key = lambda x:x['to'])

    with open(os.path.join(project_path, 'imports.tf'), 'w') as imports_file:
        for import_ in imports:
            id = import_['id'].replace('"', '\\"')
            imports_file.write(
                f'''
import {{
    to = {import_['to']}
    id = "{id}"
}}
''')


def write_removed_file(project_path, removals):
    removals = sorted(removals)

    with open(os.path.join(project_path, 'removed.tf'), 'w') as removed_file:
        for removal in removals:
            removed_file.write(
                f'''
removed {{
  from = {removal}
  
  lifecycle {{
    destroy = false
  }}
}}
''')


if __name__ == '__main__':
    main()