#!/usr/bin/env python3
import argparse
import json
import subprocess
import yaml

from typing import Dict, List, Tuple


PUNCHER_RULE_URL_TEMPLATE = "https://puncher.yandex-team.ru/?id={}"


def read_expected_rules(expected_rules_file):
    with open(expected_rules_file, "r") as stream:
        return yaml.safe_load(stream)


def get_puncher_command_for_specific_source_search(puncher_binary: str, source: str, destination: str, ports: str):
    return (
        f"{puncher_binary} get --exact --json "
        f"--source {source} "
        f"--destination {destination} "
        f"--ports {ports} "
    )


def get_puncher_command_for_any_source_search(puncher_binary: str, destination: str, ports: str):
    return (
        f"{puncher_binary} get --json "
        f"--source any "
        f"--destination {destination} "
        f"--ports {ports} "
    )


def search_in_existing_rules(puncher_binary: str, rule: dict):
    puncher_command = get_puncher_command_for_specific_source_search(
        puncher_binary, rule['source'], rule['destination'], rule['ports'])
    cmd = subprocess.run(puncher_command, shell=True, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    return parse_puncher_stdout(cmd.stdout)


def search_in_any_source_rules(puncher_binary: str, rule: dict):
    puncher_command = get_puncher_command_for_any_source_search(
        puncher_binary, rule['destination'], rule['ports'])
    cmd = subprocess.run(puncher_command, shell=True, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    return parse_puncher_stdout(cmd.stdout)


def parse_puncher_stdout(stdout) -> List[str]:
    puncher_rules_ids = []

    for line in stdout.decode().split("\n"):
        # Skip empty lines.
        if not line:
            continue

        puncher_rules_ids.append(json.loads(line).get("id"))

    return puncher_rules_ids


def get_puncher_rules(expected, puncher_binary, force_search_rules_with_source_any):
    existing: List[Dict] = []
    missing: List[Tuple[Dict, List[str]]] = []
    source_any: List[Dict] = []

    for rule in expected:

        puncher_rules_ids = search_in_existing_rules(puncher_binary, rule)

        if puncher_rules_ids:
            existing.append((rule, puncher_rules_ids))

            if force_search_rules_with_source_any:
                puncher_rules_ids = search_in_any_source_rules(puncher_binary, rule)
                if puncher_rules_ids:
                    source_any.append((rule, puncher_rules_ids))

            continue

        puncher_rules_ids = search_in_any_source_rules(puncher_binary, rule)
        if puncher_rules_ids:
            source_any.append((rule, puncher_rules_ids))
        else:
            missing.append(rule)

    return existing, missing, source_any


def rule_to_str(rule):
    return f"{rule['source']} -> {rule['destination']} {rule['ports']}"


def print_existing_rules(rules):
    print("* Existing rules")
    print("=============")
    for rule in rules:
        print(rule_to_str(rule[0]))
        for rule_id in rule[1]:
            print(PUNCHER_RULE_URL_TEMPLATE.format(rule_id))
        print()


def print_source_any_rules(rules):
    print("* Rules with \"source: any\"")
    print("=============")
    for rule in rules:
        print(rule_to_str(rule[0]))
        for rule_id in rule[1]:
            print(PUNCHER_RULE_URL_TEMPLATE.format(rule_id))
        print()


def print_missing_rules(rules):
    if not rules:
        print("* All rules defined in rules file are in Puncher.")
        return
    print("* Missing rules in Puncher")
    print("==============")
    for rule in rules:
        print(rule_to_str(rule))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Validates that rules defined in config file exist in Puncher. Ensure you have Puncher CLI "
                    "installed or built, and authenticated, see https://docs.yandex-team.ru/puncher/cli .")
    parser.add_argument("-r", "--expected-rules-file", type=str, required=True, help="YAML file with expected rules.")
    parser.add_argument("-p", "--puncher-binary", type=str, required=True,
                        help="Path to Puncher binary. Code path in Arcadia: noc/puncher/cli/puncher/puncher")
    parser.add_argument("-a", "--force-search-rules-with-source-any", action="store_true",
                        help="Search for rules with \"source: any\" even when rules with specified source are found.")
    args = parser.parse_args()

    expected_rules = read_expected_rules(args.expected_rules_file)
    existing_rules, missing_rules, source_any = get_puncher_rules(
        expected_rules, args.puncher_binary, args.force_search_rules_with_source_any)
    print_existing_rules(existing_rules)
    print()
    print_source_any_rules(source_any)
    print()
    print_missing_rules(missing_rules)
