#!/usr/bin/env python3
"""Move legacy per-user deSEC zones into the parent zone's CNAME wildcard model.

Run from the backend environment. The default is a read-only dry run; pass
``--apply`` to copy records, cut over delegations, export the old zonefiles,
and remove the child zones that would otherwise override the parent records.
"""

import argparse
import os
import re
import sys
import time
from typing import Dict, Iterable, List, Tuple
from urllib.parse import quote

import requests


def _load_usernames() -> List[str]:
    sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    from database_service import database_connection

    with database_connection() as connection:
        cursor = connection.cursor()
        cursor.execute(
            """
            SELECT DISTINCT u.username
            FROM users u
            JOIN user_domains d ON d.user_id = u.id
            WHERE COALESCE(d.is_custom, FALSE) = FALSE
              AND d.full_domain IS NOT NULL
            """
        )
        return [row['username'] for row in cursor.fetchall()]


def _valid_namespace(username: str, base_domain: str) -> str:
    namespace = f"u-{username.strip().lower()}"
    labels = namespace.split('.')
    pattern = re.compile(r"[a-z0-9](?:[a-z0-9-]*[a-z0-9])?\Z")
    if not all(len(label) <= 63 and pattern.fullmatch(label) for label in labels):
        raise ValueError(f"Username cannot form a DNS namespace: {username!r}")
    if len(f"{namespace}.{base_domain}") > 253:
        raise ValueError(f"Username namespace is too long: {username!r}")
    return namespace


class DeSec:
    def __init__(self, token: str, base_domain: str):
        self.base_url = 'https://desec.io/api/v1'
        self.base_domain = base_domain.rstrip('.').lower()
        self.headers = {'Authorization': f'Token {token}', 'Content-Type': 'application/json'}
        self._last_request = 0.0

    def request(self, method: str, url: str, **kwargs) -> requests.Response:
        for attempt in range(6):
            delay = 0.2 - (time.monotonic() - self._last_request)
            if delay > 0:
                time.sleep(delay)
            response = requests.request(method, url, headers=self.headers, timeout=30, **kwargs)
            self._last_request = time.monotonic()
            if response.status_code != 429:
                return response
            retry_after = response.headers.get('Retry-After', '2')
            try:
                wait = min(max(float(retry_after), 1), 30)
            except ValueError:
                wait = 2
            if attempt == 5:
                return response
            time.sleep(wait)
        raise RuntimeError('DNS API retry loop ended unexpectedly')

    def list_rrsets(self, zone: str) -> List[Dict]:
        url = f'{self.base_url}/domains/{zone}/rrsets/'
        rrsets: List[Dict] = []
        while url:
            response = self.request('GET', url)
            if response.status_code != 200:
                raise RuntimeError(f'Could not read RRsets for {zone} ({response.status_code})')
            result = response.json()
            if not isinstance(result, list):
                raise RuntimeError(f'Unexpected RRset response for {zone}')
            rrsets.extend(result)
            url = response.links.get('next', {}).get('url')
        return rrsets

    def rrset_url(self, subname: str, record_type: str) -> str:
        return (
            f'{self.base_url}/domains/{self.base_domain}/rrsets/'
            f'{quote(subname, safe=".-_*.")}/{record_type}/'
        )


def _rrset_records(rrset: Dict) -> Tuple[str, ...]:
    return tuple(sorted(str(record).rstrip('.').lower() if rrset.get('type') == 'CNAME' else str(record)
                         for record in rrset.get('records', [])))


def _relative_owner(zone: str, base_domain: str) -> str:
    suffix = f'.{base_domain}'
    if not zone.endswith(suffix):
        raise ValueError(f'Zone is outside configured parent domain: {zone}')
    owner = zone[:-len(suffix)]
    _validate_owner(owner)
    return owner


def _validate_owner(owner: str) -> None:
    labels = owner.split('.')
    pattern = re.compile(r"(?:\*|[a-z0-9](?:[a-z0-9-]*[a-z0-9])?)\Z")
    if (
        not all(len(label) <= 63 and pattern.fullmatch(label) for label in labels)
        or ('*' in labels and (labels[0] != '*' or labels.count('*') != 1))
    ):
        raise ValueError(f'Unsupported legacy namespace owner: {owner!r}')


def build_plan(api: DeSec) -> Tuple[List[str], List[str], List[Dict], Dict[Tuple[str, str], Dict]]:
    zones_response = api.request('GET', f'{api.base_url}/domains/')
    if zones_response.status_code != 200:
        raise RuntimeError(f'Could not list deSEC zones ({zones_response.status_code})')
    zones = zones_response.json()
    legacy_zones = sorted(
        zone['name'] for zone in zones
        if zone.get('name', '').endswith(f'.{api.base_domain}')
        and zone.get('name', '').split('.', 1)[0].startswith('u-')
    )

    parent_rrsets = api.list_rrsets(api.base_domain)
    parent_map = {(item.get('subname') or '', item['type']): item for item in parent_rrsets}
    namespaces = {_relative_owner(zone, api.base_domain) for zone in legacy_zones}
    namespaces.update(
        item.get('subname', '') for item in parent_rrsets
        if item.get('type') in {'NS', 'DS'} and (item.get('subname') or '').startswith('u-')
    )
    for username in _load_usernames():
        namespaces.add(_valid_namespace(username, api.base_domain))

    copies: List[Dict] = []
    for zone in legacy_zones:
        namespace = _relative_owner(zone, api.base_domain)
        for rrset in api.list_rrsets(zone):
            child_name = rrset.get('subname') or ''
            record_type = rrset['type']
            if child_name in {'', '@'} and record_type == 'NS':
                continue
            if child_name in {'', '@'}:
                raise RuntimeError(
                    f'Cannot preserve apex {record_type} at {zone}; manual DNS migration is required'
                )
            target_name = f'{child_name}.{namespace}'
            _validate_owner(target_name)
            key = (target_name, record_type)
            existing = parent_map.get(key)
            other_types = {
                item['type'] for item in parent_rrsets
                if item.get('subname', '') == target_name and item['type'] != record_type
            }
            if record_type == 'CNAME' and other_types:
                raise RuntimeError(f'Conflicting parent RRset at {target_name}: {sorted(other_types)}')
            if 'CNAME' in other_types:
                raise RuntimeError(f'Parent CNAME conflicts with child RRset at {target_name} {record_type}')
            if existing and _rrset_records(existing) != _rrset_records(rrset):
                raise RuntimeError(f'Conflicting parent RRset at {target_name} {record_type}')
            if not existing:
                copies.append({
                    'subname': target_name,
                    'type': record_type,
                    'ttl': rrset['ttl'],
                    'records': rrset.get('records', []),
                })

    for namespace in sorted(namespaces):
        for owner in (namespace, f'*.{namespace}'):
            _validate_owner(owner)
            current_at_owner = [item for item in parent_rrsets if item.get('subname', '') == owner]
            for item in current_at_owner:
                if item['type'] in {'NS', 'DS'} and owner == namespace:
                    continue
                if item['type'] == 'CNAME' and _rrset_records(item) == (api.base_domain,):
                    continue
                raise RuntimeError(f'Conflicting parent RRset at {owner} {item["type"]}')

    return sorted(namespaces), legacy_zones, copies, parent_map


def _delete_rrset(api: DeSec, subname: str, record_type: str) -> None:
    response = api.request('DELETE', api.rrset_url(subname, record_type))
    if response.status_code not in {200, 204, 404}:
        raise RuntimeError(f'Could not remove parent {record_type} at {subname} ({response.status_code})')


def _create_rrset(api: DeSec, payload: Dict) -> None:
    response = api.request('POST', f'{api.base_url}/domains/{api.base_domain}/rrsets/', json=payload)
    if response.status_code not in {200, 201}:
        raise RuntimeError(
            f'Could not create parent {payload["type"]} at {payload["subname"]} ({response.status_code})'
        )


def _backup_and_remove_legacy_zones(api: DeSec, legacy_zones: Iterable[str], backup_dir: str) -> None:
    """Export old zones for rollback, then remove their competing authority."""
    os.makedirs(backup_dir, mode=0o700, exist_ok=True)
    zonefiles = {}
    for zone in legacy_zones:
        url = f'{api.base_url}/domains/{quote(zone, safe=".-")}/zonefile/'
        response = api.request('GET', url)
        if response.status_code != 200:
            raise RuntimeError(f'Could not export legacy zone {zone} ({response.status_code})')
        zonefiles[zone] = response.text

    for zone, zonefile in zonefiles.items():
        filename = os.path.join(backup_dir, f'{zone}.zone')
        with open(filename, 'w', encoding='utf-8') as backup:
            backup.write(zonefile)
    print(f'Exported {len(zonefiles)} legacy zone backups to {backup_dir}')

    for zone in legacy_zones:
        url = f'{api.base_url}/domains/{quote(zone, safe=".-")}/'
        response = api.request('DELETE', url)
        if response.status_code not in {204, 404}:
            raise RuntimeError(f'Could not remove legacy zone {zone} ({response.status_code})')
        print(f'Removed legacy child zone {zone}')


def apply_plan(api: DeSec, namespaces: Iterable[str], copies: List[Dict],
               legacy_zones: Iterable[str], parent_map: Dict[Tuple[str, str], Dict],
               backup_dir: str) -> None:
    # Copy child records and create wildcard records first; old delegations stay
    # authoritative until all non-destructive preparation has completed.
    for payload in copies:
        _create_rrset(api, payload)
        print(f'Copied {payload["type"]} {payload["subname"]}')

    for namespace in namespaces:
        wildcard = f'*.{namespace}'
        if (wildcard, 'CNAME') not in parent_map:
            _create_rrset(api, {
                'subname': wildcard, 'type': 'CNAME', 'ttl': 3600,
                'records': [f'{api.base_domain}.'],
            })
            print(f'Created wildcard CNAME {wildcard}')

    for namespace in namespaces:
        # Remove DS before NS so a resolver never sees a signed delegation with
        # no matching child zone. Preserve exact RRsets for a one-step rollback.
        saved = {record_type: parent_map.get((namespace, record_type))
                 for record_type in ('DS', 'NS')}
        try:
            for record_type in ('DS', 'NS'):
                if saved[record_type]:
                    _delete_rrset(api, namespace, record_type)
            if (namespace, 'CNAME') not in parent_map:
                _create_rrset(api, {
                    'subname': namespace, 'type': 'CNAME', 'ttl': 3600,
                    'records': [f'{api.base_domain}.'],
                })
            print(f'Cut over {namespace}')
        except Exception:
            # The old child zones remain intact. Restore delegation metadata if
            # the final CNAME write fails partway through a cutover.
            for record_type in ('NS', 'DS'):
                rrset = saved[record_type]
                if rrset:
                    _create_rrset(api, {
                        'subname': namespace,
                        'type': record_type,
                        'ttl': rrset['ttl'],
                        'records': rrset['records'],
                    })
            raise

    # Hosted child zones continue to answer authoritatively even after their
    # parent NS/DS records are removed. Back them up and remove them so the
    # new parent-zone CNAMEs and wildcards are actually served.
    _backup_and_remove_legacy_zones(api, legacy_zones, backup_dir)

    from database_service import database_connection

    with database_connection() as connection:
        cursor = connection.cursor()
        for namespace in namespaces:
            pattern = f'%.{namespace}.{api.base_domain}'
            pattern = pattern.replace('!', '!!').replace('%', '!%').replace('_', '!_')
            cursor.execute(
                """
                UPDATE user_domains
                   SET record_type = ?, ip_address = ?, ttl = 3600
                 WHERE COALESCE(is_custom, FALSE) = FALSE
                   AND (lower(full_domain) = ? OR lower(full_domain) LIKE ? ESCAPE '!')
                """,
                ('CNAME', api.base_domain, f'{namespace}.{api.base_domain}', pattern),
            )
        connection.commit()


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--apply', action='store_true', help='Apply the reviewed DNS migration')
    parser.add_argument(
        '--backup-dir', default='/tmp/anchorhost-delegated-dns-backup',
        help='Directory for exported legacy zonefiles before their zones are deleted',
    )
    args = parser.parse_args()
    token = os.environ.get('DESEC_API_TOKEN')
    base_domain = os.environ.get('BASE_DOMAIN')
    if not token or not base_domain:
        parser.error('DESEC_API_TOKEN and BASE_DOMAIN must be set')

    api = DeSec(token, base_domain)
    try:
        namespaces, legacy_zones, copies, parent_map = build_plan(api)
        print(f'Namespaces to migrate: {len(namespaces)}')
        print(f'Legacy RRsets to copy: {len(copies)}')
        print('Namespaces:', ', '.join(namespaces))
        print('Record copies:', ', '.join(f'{item["type"]}:{item["subname"]}' for item in copies) or 'none')
        print('Mode:', 'APPLY' if args.apply else 'DRY RUN')
        if args.apply:
            apply_plan(api, namespaces, copies, legacy_zones, parent_map, args.backup_dir)
            print('DNS migration completed. Legacy zonefiles are backed up; old child zones are removed.')
        else:
            print('No changes made. Run with --apply to perform the migration.')
    except Exception as error:
        print(f'Migration stopped: {error}', file=sys.stderr)
        return 1
    return 0


if __name__ == '__main__':
    raise SystemExit(main())