#!/usr/bin/env python3
import yaml
import json
import logging
import argparse
import requests
from os.path import dirname, abspath

parser = argparse.ArgumentParser()
parser.add_argument('--mode', type=str, required=True, choices=["deploy", "check", "mock"])
opt = parser.parse_args()
base_path = dirname(dirname(abspath(__file__)))
schemas_conf = yaml.safe_load(open(f'{base_path}/schemas.yml', 'r'))
clusters = yaml.safe_load(open(f'{base_path}/clusters.yml', 'r'))
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

headers = {
    'Content-type': 'application/json',
    'Accept': 'application/vnd.schemaregistry.v1+json'
}

deployment_envs = ['devc', 'staging', 'prod']

registered_topics = {}

with requests.Session() as request_session:
    request_session.headers.update(headers)


def _post_schema(registry_base_url: str, topic_name: str, filepath: str) -> int:
    with open(f'{base_path}/schemas/{filepath}', 'r') as schemas_conf_file:
        schema_str = schemas_conf_file.read()
    schema_definition = {'schema': schema_str, 'schemaType': 'PROTOBUF'}
    response = request_session.post(url=f'{registry_base_url}/subjects/{topic_name}-value/versions',
                                    data=json.dumps(schema_definition))
    response.raise_for_status()
    return response.json()["id"]


def _register_id(topic_name: str, deploy_env: str, deploy_region: str, deploy_az: str, schema_id: int) -> None:
    envs = registered_topics.get(topic_name, {})
    regions = envs.get(deploy_env, {})
    ids = regions.get(deploy_region, {})
    ids[deploy_az] = schema_id
    regions[deploy_region] = ids
    envs[deploy_env] = regions
    registered_topics[topic_name] = envs


def deploy_schema(deploy_env: str, deploy_region: str, deploy_az: str, registry_base_url: str) -> dict:
    clusters_schemas_ids = {}
    for topic_name, schema_conf in schemas_conf.items():
        filepath = schema_conf.get('filepath')
        logger.info(f'registering schema for topic : {topic_name} in {region}{az}')
        schema_id = _post_schema(registry_base_url, topic_name, filepath)
        _register_id(topic_name, deploy_env, deploy_region, deploy_az, schema_id)
    return clusters_schemas_ids


def check_schema_id(registry_base_url: str, topic_name: str, expected_id: int) -> None:
    logger.info(f'checking schema id for topic : {topic_name} on registry_base_url')
    filepath = schemas_conf[topic_name]['filepath']
    schema_id = _post_schema(registry_base_url, topic_name, filepath)
    assert schema_id == expected_id


if opt.mode == "deploy":
    env_schemas = {}
    for env in deployment_envs:
        regions_conf = clusters.get(env)
        for region, az_conf in regions_conf.items():
            for az, url in az_conf.items():
                try:
                    requests.put(url=f'{url}/mode', data=json.dumps({'mode': 'READWRITE'}))
                    deploy_schema(env, region, az, url)
                finally:
                    requests.put(url=f'{url}/mode', data=json.dumps({'mode': 'READONLY'}))
    for topic, topic_ids in registered_topics.items():
        with open(f'{base_path}/ids/{topic}.json', 'w+') as env_ids_file:
            env_ids_file.write(json.dumps(topic_ids, indent=4))
elif opt.mode == "check":
    for topic in schemas_conf:
        with open(f'{base_path}/ids/{topic}.json', 'r') as topic_ids_json:
            topic_ids = json.load(topic_ids_json)
        for env in deployment_envs:
            for region, az_conf in clusters.get(env).items():
                for az, url in az_conf.items():
                    topic_id = topic_ids[env][region][az]
                    check_schema_id(url, topic, topic_id)
elif opt.mode == 'mock':
    local_env = 'local'
    for region, az_conf in clusters.get(local_env).items():
        for az, url in az_conf.items():
            az_schemas = deploy_schema(local_env, region, az, url)

request_session.close()