import json
import sys


def json_to_mermaid(json_file_path, output_mmd_path='network_diagram.mmd'):
    """
    Converts the network discovery JSON to a Mermaid flowchart.
    """
    with open(json_file_path, 'r', encoding='utf-8') as f:
        data = json.load(f)

    mermaid_lines = []
    mermaid_lines.append('flowchart TD')

    # 1. Define styling classes
    mermaid_lines.append('    %% ---------- styles ----------')
    mermaid_lines.append('    classDef internet fill:#e1f5ff,stroke:#007bff')
    mermaid_lines.append('    classDef router  fill:#fff3cd,stroke:#ffc107')
    mermaid_lines.append('    classDef server  fill:#ffe6e6,stroke:#dc3545')
    mermaid_lines.append('    classDef dns     fill:#e6ffe6,stroke:#28a745')
    mermaid_lines.append('    classDef homeauto fill:#fff0f5,stroke:#e83e8c')
    mermaid_lines.append('    classDef worker  fill:#f0e6ff,stroke:#6f42c1')
    mermaid_lines.append('    classDef iot     fill:#fff9e6,stroke:#fd7e14')
    mermaid_lines.append('    classDef unknown fill:#f8f9fa,stroke:#6c757d')
    mermaid_lines.append('')

    # 2. Create a mapping for device type to CSS class
    type_to_class = {
        'router': 'router',
        'server': 'server',
        'dns': 'dns',
        'home_automation': 'homeauto',
        'worker': 'worker',
        'iot': 'iot'
    }

    # 3. Group devices by network
    lan_devices = []
    tether_devices = []
    hyperv_devices = []
    other_devices = []

    for device in data['devices']:
        ip = device.get('ip', '')
        if ip.startswith('192.168.1.'):
            lan_devices.append(device)
        elif ip.startswith('192.168.137.'):
            tether_devices.append(device)
        elif ip.startswith('172.20.'):
            hyperv_devices.append(device)
        else:
            other_devices.append(device)

    # 4. Create LAN Subgraph
    mermaid_lines.append(f'    subgraph LAN["LAN 192.168.1.0/24 - {len(lan_devices)} devices"]')

    # Core infrastructure first (from your original diagram)
    core_ips = ['192.168.1.1', '192.168.1.159', '192.168.1.163',
                '192.168.1.193', '192.168.1.100', '192.168.1.49', '192.168.1.165']

    # Add a core server subgraph
    mermaid_lines.append('        subgraph CORE["Core Infrastructure"]')
    for device in lan_devices:
        if device['ip'] in core_ips:
            node_id = device['ip'].replace('.', '_')
            # Use ASCII-friendly labels instead of emoji for Windows compatibility
            icon_map = {
                'router': '[Router]',
                'server': '[Server]',
                'dns': '[DNS]',
                'home_automation': '[Home]',
                'worker': '[Worker]',
                'unknown': '[?]'
            }
            icon = icon_map.get(device.get('type', 'unknown'), '[?]')
            label = f"{icon} {device.get('hostname', device['ip'])}"
            mermaid_lines.append(f'            {node_id}["{label}<br/>{device["ip"]}"]')

            # Apply CSS class
            dev_type = device.get('type', 'unknown')
            css_class = type_to_class.get(dev_type, 'unknown')
            mermaid_lines.append(f'            class {node_id} {css_class};')
    mermaid_lines.append('        end')

    # Other LAN devices
    for device in lan_devices:
        if device['ip'] not in core_ips:
            node_id = device['ip'].replace('.', '_')
            icon_map = {
                'router': '[R]',
                'server': '[S]',
                'dns': '[D]',
                'home_automation': '[H]',
                'worker': '[W]',
                'unknown': '[?]'
            }
            icon = icon_map.get(device.get('type', 'unknown'), '[?]')
            label = f"{icon} {device.get('hostname', device['ip'])}"
            mermaid_lines.append(f'        {node_id}["{label}<br/>{device["ip"]}"]')

            dev_type = device.get('type', 'unknown')
            css_class = type_to_class.get(dev_type, 'unknown')
            mermaid_lines.append(f'        class {node_id} {css_class};')

    mermaid_lines.append('    end')
    mermaid_lines.append('')

    # 5. Create TETHER Subgraph
    if tether_devices:
        mermaid_lines.append(f'    subgraph TETHER["Tether 192.168.137.0/24 - {len(tether_devices)} devices"]')

        # Known workers first
        known_workers = ['192.168.137.239']  # Hermes
        for device in tether_devices:
            node_id = device['ip'].replace('.', '_')
            if device['ip'] in known_workers:
                label = f"[Hermes] {device.get('hostname', device['ip'])}<br/>{device['ip']}"
            else:
                label = f"[Worker] {device.get('hostname', device['ip'])}<br/>{device['ip']}"
            mermaid_lines.append(f'        {node_id}["{label}"]')
            mermaid_lines.append(f'        class {node_id} worker;')

        mermaid_lines.append('    end')
        mermaid_lines.append('')

    # 6. Add connections
    mermaid_lines.append('    %% ---------- Connections ----------')
    # Connect core server to tether network
    if tether_devices and any(d['ip'] == '192.168.1.159' for d in lan_devices):
        mermaid_lines.append('    192_168_1_159 -.-> TETHER')
    # Router connects to everything
    if any(d['ip'] == '192.168.1.1' for d in lan_devices):
        mermaid_lines.append('    192_168_1_1 --> LAN')

    # 7. Write the Mermaid code to a file with UTF-8 encoding
    with open(output_mmd_path, 'w', encoding='utf-8') as f:
        f.write('\n'.join(mermaid_lines))

    print(f"✓ Mermaid diagram generated: {output_mmd_path}")
    print(f"  Devices: {len(lan_devices)} LAN, {len(tether_devices)} Tether")
    if hyperv_devices:
        print(f"           {len(hyperv_devices)} Hyper-V")
    print(f"  Open {output_mmd_path} and copy contents to https://mermaid.live")
    return output_mmd_path


# --- Run the conversion with command-line support ---
if __name__ == "__main__":
    # Default filenames
    input_json = "network_discovery_20251207_160645.json"
    output_mmd = "network_topology.mmd"

    # Check for command-line arguments
    if len(sys.argv) > 1:
        input_json = sys.argv[1]
        if len(sys.argv) > 2:
            output_mmd = sys.argv[2]

    try:
        mmd_file = json_to_mermaid(input_json, output_mmd)
    except FileNotFoundError:
        print(f"✗ Error: Could not find file '{input_json}'")
        print("  Available JSON files:")
        import os

        for file in os.listdir('.'):
            if file.endswith('.json'):
                print(f"    - {file}")
    except json.JSONDecodeError as e:
        print(f"✗ Error: Invalid JSON in '{input_json}': {e}")
    except Exception as e:
        print(f"✗ Unexpected error: {e}")
        import traceback

        traceback.print_exc()