import grpc
import json
import requests
from common.proto.context_pb2 import DeviceId, DeviceOperationalStatusEnum, ConfigRule, ConfigRule_Custom, Empty, ContextId, TopologyId, Uuid, ConfigActionEnum
from common.proto.context_pb2_grpc import ContextServiceStub
from common.DeviceTypes import DeviceTypeEnum
from common.tools.grpc.ConfigRules import update_config_rule_custom

def login_qkd(address, port, username, password):
    url = f"http://{address}:{port}/login"
    headers = {'Content-Type': 'application/x-www-form-urlencoded'}
    data = {'username': username, 'password': password}
    response = requests.post(url, headers=headers, data=data)
    response.raise_for_status()
    token = response.json().get('access_token')
    return token

def fetch_qkd_info(address, port, token, endpoint):
    url = f"http://{address}:{port}{endpoint}"
    headers = {'Authorization': f'Bearer {token}', 'accept': 'application/json'}
    response = requests.get(url, headers=headers)
    response.raise_for_status()
    return response.json()

def create_config_file(qkd_devices, filename):
    json_structure = {
        "contexts": [
            {"context_id": {"context_uuid": {"uuid": "admin"}}}
        ],
        "topologies": [
            {"topology_id": {"topology_uuid": {"uuid": "admin"}, "context_id": {"context_uuid": {"uuid": "admin"}}}}
        ],
        "devices": [],
        "links": []
    }

    for device in qkd_devices:
        device_entry = {
            "device_id": {"device_uuid": {"uuid": device["uuid"]}},
            "device_type": "qkd-node",
            "device_operational_status": DeviceOperationalStatusEnum.DEVICEOPERATIONALSTATUS_ENABLED,
            "device_drivers": [12],  # Assuming 12 is the correct driver ID for QKD
            "device_endpoints": [],
            "device_config": {"config_rules": [
                {"action": ConfigActionEnum.CONFIGACTION_SET, "custom": {"resource_key": "_connect/address", "resource_value": device["address"]}},
                {"action": ConfigActionEnum.CONFIGACTION_SET, "custom": {"resource_key": "_connect/port", "resource_value": str(device["port"])}},
                {"action": ConfigActionEnum.CONFIGACTION_SET, "custom": {"resource_key": "_connect/settings", "resource_value": json.dumps({"scheme": "http", "token": device["token"]})}}
            ]}
        }
        json_structure["devices"].append(device_entry)

        for interface in device["interfaces"]["qkd_interface"]:
            endpoint_id = f"{device['address']}:{interface['qkdi_id']}"
            device_entry["device_endpoints"].append({
                "device_id": {"device_uuid": {"uuid": device["uuid"]}},
                "endpoint_uuid": {"uuid": endpoint_id}
            })

        for link in device["links"]["qkd_links"]:
            link_entry = {
                "link_id": {"link_uuid": {"uuid": link["qkdl_id"]}},
                "link_endpoint_ids": [
                    {"device_id": {"device_uuid": {"uuid": device["uuid"]}}, "endpoint_uuid": {"uuid": f"{device['address']}:{link['qkdl_id']}"}}
                    # You need to fetch and add the other endpoint details similarly
                ]
            }
            json_structure["links"].append(link_entry)

    with open(filename, 'w', encoding='utf-8') as f:
        json.dump(json_structure, f, ensure_ascii=False, indent=4)

def main():
    qkd_devices = [
        {
            "uuid": "real_qkd1",
            "address": "10.13.13.2",
            "port": 5100,
            "username": "admin",
            "password": "password"
        }
        # Add more QKD devices if needed
    ]

    # Step 1: Authenticate and get the JWT tokens for each QKD device
    for qkd_device in qkd_devices:
        qkd_device['token'] = login_qkd(
            qkd_device['address'], qkd_device['port'], qkd_device['username'], qkd_device['password'])

    # Step 2: Fetch QKD device information
    for qkd_device in qkd_devices:
        qkd_device['capabilities'] = fetch_qkd_info(qkd_device['address'], qkd_device['port'], qkd_device['token'], '/restconf/data/etsi-qkd-sdn-node:qkd_node/qkdn_capabilities')
        qkd_device['interfaces'] = fetch_qkd_info(qkd_device['address'], qkd_device['port'], qkd_device['token'], '/restconf/data/etsi-qkd-sdn-node:qkd_node/qkd_interfaces')
        qkd_device['links'] = fetch_qkd_info(qkd_device['address'], qkd_device['port'], qkd_device['token'], '/restconf/data/etsi-qkd-sdn-node:qkd_node/qkd_links')
        qkd_device['applications'] = fetch_qkd_info(qkd_device['address'], qkd_device['port'], qkd_device['token'], '/restconf/data/etsi-qkd-sdn-node:qkd_node/qkd_applications')
    
    # Step 3: Create config files for each QKD device
    config_filename = "qkd_devices_config.json"
    create_config_file(qkd_devices, config_filename)
    print(f"Config file created: {config_filename}")

if __name__ == "__main__":
    main()