from typing import NamedTuple
import re
from io import StringIO
import sys
import tomllib
from pathlib import Path
from typing import TypedDict
from collections import defaultdict
from urllib.request import urlopen
import argparse
import json
import subprocess

TRAEFIK_RULE_RE = re.compile(r"([a-zA-Z]+)\(`(.+?)`\)")
SUBDOMAIN_WILDCARD_PREFIX = "{subdomain:[a-z]+}"


class Route(NamedTuple):
    name: str
    destination: str
    hostname: str


class Source(TypedDict):
    type: str
    url: str
    destination: str


class Config(TypedDict):
    source: list[Source]


def parse_traefik_rule(rule: str) -> list[str]:
    rules = defaultdict(list)
    for (rule_fn, rule_arg) in TRAEFIK_RULE_RE.findall(rule):
        rules[rule_fn].append(rule_arg)

    if host_rules := rules.get("Host"):
        return host_rules[0]

    elif regex_rules := rules.get("HostRegexp"):
        regex_rule = regex_rules[0]
        if regex_rule.startswith(f"{SUBDOMAIN_WILDCARD_PREFIX}."):
            return regex_rule.replace(SUBDOMAIN_WILDCARD_PREFIX, "*", 1)

    return None


def get_traefik_routes(traefik_host: str, traefik_route: str):
    with urlopen(f"{traefik_host}/api/http/routers") as response:
        api_response = json.load(response)

    routes = set()

    for router in api_response:
        host = parse_traefik_rule(router["rule"])

        if not host:
            print(f"Failed to find host for {router['rule']}", file=sys.stderr)
            continue

        routes.add(Route(
            router["service"],
            traefik_route,
            host
        ))

    return routes

def get_dokku_routes(dokku_exporter_url: str, dokku_route: str):
    with urlopen(dokku_exporter_url) as response:
        api_response = json.load(response)

    routes = set()

    for app in api_response:
        for vhost in app["vhosts"]:
            routes.add(Route(
                app["app"],
                dokku_route,
                vhost
            ))

    return routes


def main():
    parser = argparse.ArgumentParser()

    parser.add_argument("--config", "-c", default="./config.toml", type=Path, help="Config file. Default: %(default)s")
    parser.add_argument("--output", "-o", type=Path, default="./map.txt", help="Output file. Default: %(default)s")
    parser.add_argument("--reload", action="store_true", help="Reload nginx on completion")

    args = parser.parse_args()

    config: Config = tomllib.loads(args.config.read_text())

    routes = []
    for source in config["source"]:
        match source["type"]:
            case "traefik":
                routes.extend(get_traefik_routes(source["url"], source["destination"]))
            case "dokku":
                routes.extend(get_dokku_routes(source["url"], source["destination"]))

    if len(routes) != len(set(routes)):
        raise ValueError("Conflict found!")

    grouped_routes = defaultdict(set)
    for route in routes:
        grouped_routes[route.hostname, route.destination].add(route.name)

    print("Found", len(routes), "routes, grouped into", len(grouped_routes), "groups")

    output = StringIO()
    for (hostname, destination), names in sorted(grouped_routes.items()):
        output.write(f"{hostname}\t{destination};  # {', '.join(names)}\n")

    if args.output.is_file():
        current_output = args.output.read_text()

        if current_output == output.getvalue():
            print("Output identical - aborting")
            return

    args.output.write_text(output.getvalue())

    if args.reload:
        print("Reloading nginx")
        subprocess.check_call(
            [
                "nginx",
                "-s",
                "reload"
            ],
        )

if __name__ == "__main__":
    main()