"""DNS graph: the delegation path from the root down to the domain, and what
the domain points at (NS, MX, web names, CNAME chains, addresses, networks).

Output is a plain node/edge list with a `layer` per node; the GUI lays it out
in columns, left to right."""

from concurrent.futures import ThreadPoolExecutor

from . import dnsutil
from .infra import asn

_pool = ThreadPoolExecutor(max_workers=24, thread_name_prefix='gr')

def build(domain):
    nodes, edges = {}, []

    def node(nid, label, kind, layer, **extra):
        if nid not in nodes:
            nodes[nid] = {'id': nid, 'label': label, 'kind': kind, 'layer': layer, **extra}
        return nid

    def edge(a, b, label=None, **extra):
        e = {'from': a, 'to': b, 'label': label, **extra}
        if e not in edges:
            edges.append(e)

    steps = dnsutil.delegation_chain(domain, 'SOA')
    # Delegation path: one column per zone cut, root first.
    prev = node('zone:.', '.', 'root', 0)
    depth = 0
    for st in steps:
        if 'child' not in st:
            continue
        depth += 1
        child = st['child']
        kind = 'tld' if is_tld(child) else 'zone'
        cid = node('zone:' + child, child, kind, depth, ns=st['ns'], signed=st['ds'],
                   ms=st.get('ms'), answered_by=st.get('server'))
        edge(prev, cid, 'NS', signed=st['ds'])
        prev = cid
    last = steps[-1] if steps else {}
    zone_node = prev
    if nodes[zone_node]['label'] != domain:
        # The name is not its own zone (e.g. www.example.com inside
        # example.com): hang it off the closest enclosing zone.
        depth += 1
        zone_node = node('zone:' + domain, domain, 'name', depth)
        edge(prev, zone_node, None)
    if last.get('error'):
        nodes[zone_node]['error'] = last['error']
    if last.get('rcode') == 'NXDOMAIN':
        nodes[zone_node]['nxdomain'] = True
        return {'domain': domain, 'nodes': list(nodes.values()), 'edges': edges,
                'delegation': steps}
    # Columns after the zone: names it publishes (NS, MX, www), CNAME targets,
    # addresses, networks. Empty columns are collapsed by the GUI.
    L_REC, L_HOST, L_IP, L_NET = depth + 1, depth + 2, depth + 3, depth + 4

    ns = [n.rstrip('.') for n in dnsutil.resolve(domain, 'NS')]
    mx = sorted((v.split(None, 1) for v in dnsutil.resolve(domain, 'MX')), key=lambda p: int(p[0]))
    web = [domain, 'www.' + domain]

    host_ips = {}

    def addrs(h):
        return h, [i for i in dnsutil.addresses(h) if dnsutil.is_public_ip(i)][:4]

    hosts = ns[:8] + [h.rstrip('.') for _, h in mx[:6] if h not in ('.', '')]
    for h, ips in _pool.map(addrs, hosts):
        host_ips[h] = ips

    for n in ns[:8]:
        hid = node('host:' + n, n, 'ns', L_REC)
        edge(zone_node, hid, 'NS')
    for pref, h in mx[:6]:
        h = h.rstrip('.')
        if h in ('', '.'):
            continue
        hid = node('host:' + h, h, 'mx', L_REC, pref=int(pref))
        edge(zone_node, hid, f'MX {pref}')

    chains = dict(_pool.map(lambda w: (w, dnsutil.cname_chain(w)), web))
    web_ips = dict(_pool.map(addrs, web))
    for w in web:
        if not chains[w] and not web_ips[w]:
            continue
        wid = zone_node if w == domain else node('name:' + w, w, 'name', L_REC)
        if w != domain:
            edge(zone_node, wid, None)
        cur = wid
        for i, target in enumerate(chains[w]):
            tid = node('host:' + target, target, 'cname', L_HOST)
            edge(cur, tid, 'CNAME')
            cur = tid
        if chains[w]:
            host_ips[chains[w][-1]] = web_ips[w]
        else:
            for ip in web_ips[w]:
                iid = node('ip:' + ip, ip, 'ip', L_IP)
                edge(cur, iid, 'AAAA' if ':' in ip else 'A')

    for h, ips in host_ips.items():
        hid = 'host:' + h
        if hid not in nodes:
            continue
        for ip in ips:
            iid = node('ip:' + ip, ip, 'ip', L_IP)
            edge(hid, iid, 'AAAA' if ':' in ip else 'A')

    ip_nodes = [n for n in nodes.values() if n['kind'] == 'ip'][:24]
    for n, info in zip(ip_nodes, _pool.map(lambda n: asn(n['label']), ip_nodes)):
        if not info:
            continue
        aid = node(f"as:{info['asn']}", f"AS{info['asn']}", 'asn', L_NET,
                   name=info['name'], country=info['country'])
        edge(n['id'], aid, info['prefix'])

    return {'domain': domain, 'nodes': list(nodes.values()), 'edges': edges,
            'delegation': steps}


def is_tld(name):
    return '.' not in name.strip('.')