#!/usr/bin/env python3
"""Stage UniFi firewall groups for the current network segmentation.

Default mode is dry-run. Pass --apply to create/update/delete as needed.
Reads credentials from the runtime automation .env on PD by default.
"""

from __future__ import annotations

import argparse
import importlib.util
import json
import os
from pathlib import Path
from typing import Any

DEFAULT_RUNTIME_HELPER = Path('/mnt/docker-ssd/docker/compose/automation/bin/unifi_stage_low_risk_objects.py')
DEFAULT_ENV_PATH = Path('/mnt/docker-ssd/docker/compose/automation/.env')


def load_helper(path: Path):
    spec = importlib.util.spec_from_file_location('unifi_stage', path)
    if spec is None or spec.loader is None:
        raise RuntimeError(f'Could not load helper module from {path}')
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def desired_groups() -> list[dict[str, Any]]:
    return [
        {'name': 'NET-MGMT', 'group_type': 'address-group', 'group_members': ['10.5.0.0/24']},
        {'name': 'NET-TRUSTED', 'group_type': 'address-group', 'group_members': ['10.5.1.0/24']},
        {'name': 'NET-SERVERS', 'group_type': 'address-group', 'group_members': ['10.5.30.0/24']},
        {'name': 'NET-IOT', 'group_type': 'address-group', 'group_members': ['10.5.10.0/24']},
        {'name': 'NET-GUEST', 'group_type': 'address-group', 'group_members': ['10.5.90.0/24']},
        {'name': 'NET-CAMERAS', 'group_type': 'address-group', 'group_members': ['10.5.20.0/24']},
        {'name': 'NET-LEGACY-CIA', 'group_type': 'address-group', 'group_members': ['192.168.1.0/24']},
        {
            'name': 'NET-RFC1918-ALL',
            'group_type': 'address-group',
            'group_members': [
                '10.5.0.0/24',
                '10.5.1.0/24',
                '10.5.10.0/24',
                '10.5.20.0/24',
                '10.5.30.0/24',
                '10.5.90.0/24',
                '192.168.1.0/24',
                '192.168.2.0/24',
                '192.168.3.0/24',
            ],
        },
        {'name': 'PORT-DNS', 'group_type': 'port-group', 'group_members': ['53']},
        {'name': 'PORT-DHCP', 'group_type': 'port-group', 'group_members': ['67-68']},
        {'name': 'PORT-NTP', 'group_type': 'port-group', 'group_members': ['123']},
        {'name': 'PORT-WEB-ADMIN', 'group_type': 'port-group', 'group_members': ['80', '443', '8443', '9443']},
        {'name': 'PORT-SSH', 'group_type': 'port-group', 'group_members': ['22']},
        {'name': 'PORT-MDNS', 'group_type': 'port-group', 'group_members': ['5353']},
    ]


TEMP_GROUP_NAMES = {'TEST-TMP', 'TEST-PORT-TMP', 'TEST-PORT-TMP2', 'TEST-PORT-TMP3'}


def main() -> int:
    parser = argparse.ArgumentParser(description='Stage UniFi firewall groups.')
    parser.add_argument('--helper', default=str(DEFAULT_RUNTIME_HELPER), help='Path to runtime UniFi helper on PD')
    parser.add_argument('--env-file', default=str(DEFAULT_ENV_PATH), help='Path to runtime .env on PD')
    parser.add_argument('--site', default='default')
    parser.add_argument('--apply', action='store_true')
    args = parser.parse_args()

    helper = load_helper(Path(args.helper))
    helper.load_env_file(Path(args.env_file))
    client = helper.UniFiClient(
        base_url=os.environ.get('UNIFI_BASE_URL', 'https://10.5.0.1'),
        username=os.environ['UNIFI_USERNAME'],
        password=os.environ['UNIFI_PASSWORD'],
        verify_tls=helper.env_bool('UNIFI_VERIFY_TLS', False),
        timeout=20,
    )
    client.login()

    groups = client.get(client.site_path(args.site, 'rest/firewallgroup'))
    by_name = {g.get('name'): g for g in groups}
    result: dict[str, Any] = {'mode': 'apply' if args.apply else 'dry-run', 'changes': []}

    for group in desired_groups():
        existing = by_name.get(group['name'])
        if not existing:
            result['changes'].append({'name': group['name'], 'action': 'create', 'payload': group})
            if args.apply:
                status, payload = client._json_request('POST', client.site_path(args.site, 'rest/firewallgroup'), group)
                if status >= 300:
                    result['changes'][-1]['error'] = {'status': status, 'payload': payload}
            continue

        current_members = existing.get('group_members', []) or []
        desired_members = group.get('group_members', []) or []
        current_type = existing.get('group_type')
        if current_type != group['group_type'] or current_members != desired_members:
            updated = dict(existing)
            updated['group_type'] = group['group_type']
            updated['group_members'] = desired_members
            result['changes'].append({
                'name': group['name'],
                'action': 'update',
                'before': {'group_type': current_type, 'group_members': current_members},
                'after': {'group_type': group['group_type'], 'group_members': desired_members},
            })
            if args.apply:
                status, payload = client._json_request('PUT', client.site_path(args.site, f"rest/firewallgroup/{existing['_id']}"), updated)
                if status >= 300:
                    result['changes'][-1]['error'] = {'status': status, 'payload': payload}
        else:
            result['changes'].append({'name': group['name'], 'action': 'noop'})

    for name in sorted(TEMP_GROUP_NAMES):
        existing = by_name.get(name)
        if not existing:
            continue
        result['changes'].append({'name': name, 'action': 'delete-temp'})
        if args.apply:
            status, payload = client._json_request('DELETE', client.site_path(args.site, f"rest/firewallgroup/{existing['_id']}"))
            if status >= 300:
                result['changes'][-1]['error'] = {'status': status, 'payload': payload}

    print(json.dumps(result, indent=2, sort_keys=True))
    return 0


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