#!/usr/bin/env python3

"""
O-RU M-Plane NETCONF Server for gNB Antenna Control

Lightweight NETCONF/SSH server implementing the O-RAN WG4 M-Plane interface
for antenna mask configuration (o-ran-uplane-conf.yang).

The gNB acts as an O-RU for antenna control purposes:
  - Advertises mplane-supported-trx-control-masks (capability report)
  - Accepts tx-array-antenna-mask-config via NETCONF edit-config
  - Translates received mask to internal antenna control (localhost:9999)
  - Sends mplane-trx-control-ant-mask-update notification (logged)

Protocol:
  SSH transport (paramiko) -> NETCONF base:1.0 framing -> XML RPC processing

Architecture:
  o1-adapter (O-DU) --NETCONF:831--> mplane_server.py (O-RU) --localhost:9999--> nr_antenna_control.c
"""

import base64
import json
import logging
import os
import re
import signal
import socket
import sys
import threading
import time
import xml.etree.ElementTree as ET

import paramiko

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

MPLANE_PORT = 831
MPLANE_USER = MPLANE_PASS = "netconf"
GNB_CTRL_HOST = "127.0.0.1"
GNB_CTRL_PORT = 9999
WG4_CAPABILITY_FILE = "/opt/mplane/wg4-capability.json"
SSH_HOST_KEY = "/opt/mplane/ssh_host_rsa_key"

NS_NC = "urn:ietf:params:xml:ns:netconf:base:1.0"
NS_UPLANE = "urn:o-ran:uplane-conf:1.0"
NETCONF_END = "]]>]]>" # https://www.rfc-editor.org/rfc/rfc4742 NETCONF base:1.0 message delimiter

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(name)s: %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    stream=sys.stdout,
)
log = logging.getLogger("mplane-server")

# WG4 Capability (O-RU supported masks)
def b64_to_mask(b64_str):
    return int.from_bytes(base64.b64decode(b64_str), byteorder="big")

def mask_to_b64(mask_int):
    byte_len = max(1, (mask_int.bit_length() + 7) // 8)
    return base64.b64encode(mask_int.to_bytes(byte_len, byteorder="big")).decode()


class Capability:
    """O-RU capability report: tx-arrays with supported masks."""

    def __init__(self):
        self.tx_arrays = {}       # name -> {"masks": [{index, b64, int}], ...}
        self.current_config = {}  # array-name -> {"index": i, "mask_int": m}

    def load(self, path):
        if not (path and os.path.exists(path)):
            # Fallback: generate for 4 antennas
            masks = []
            for n in range(4, 0, -1):
                m = (1 << n) - 1
                masks.append({"index": 4 - n, "b64": mask_to_b64(m), "int": m})
            self.tx_arrays["tx-array-0"] = {"masks": masks, "columns": 4, "layers": 4}
            self.current_config["tx-array-0"] = {"index": 0, "mask_int": 0x0F}
            log.info("Generated fallback capability (4 antennas)")
            return
            
        with open(path) as f:
            data = json.load(f) # returns a dict
        for arr in data.get("tx-arrays", []): # blank list for fallback
            name = arr["name"]
            cap = arr.get("mplane-trx-control-txarr-capability-info", {})
            masks = []
            for m in cap.get("mplane-supported-trx-control-masks", []):
                masks.append({
                    "index": m["index"],
                    "b64": m["antenna-mask"],
                    "int": b64_to_mask(m["antenna-mask"]),
                })
            self.tx_arrays[name] = {
                "masks": masks,
                "columns": arr.get("number-of-columns", 4),
                "layers": arr.get("number-of-array-layers", 4),
            }
            # Default: all active (index 0)
            if masks:
                self.current_config[name] = {
                    "index": masks[0]["index"],
                    "mask_int": masks[0]["int"],
                }
        log.info("Loaded capability: %d tx-arrays from %s", len(self.tx_arrays), path)

    def resolve_index(self, array_name, index):
        arr = self.tx_arrays.get(array_name)
        if not arr:
            return None, f"unknown tx-array '{array_name}'"
        for m in arr["masks"]:
            if m["index"] == index:
                return m, ""
        return None, f"index {index} not in supported masks"

    def resolve_mask(self, array_name, mask_int):
        arr = self.tx_arrays.get(array_name)
        if not arr:
            return None, f"unknown tx-array '{array_name}'"
        for m in arr["masks"]:
            if m["int"] == mask_int:
                return m, ""
        return None, f"mask 0x{mask_int:02X} not in supported masks"

    def get_capability_xml(self):
        """Generate o-ran-uplane-conf operational data XML for <get> response."""
        parts = [f'<user-plane-configuration xmlns="{NS_UPLANE}">']
        for name, arr in self.tx_arrays.items():
            parts.append(f"<tx-arrays><name>{name}</name>")
            parts.append(f"<number-of-columns>{arr['columns']}</number-of-columns>")
            parts.append(f"<number-of-array-layers>{arr['layers']}</number-of-array-layers>")
            parts.append("<mplane-trx-control-txarr-capability-info>")
            for m in arr["masks"]:
                parts.append("<mplane-supported-trx-control-masks>")
                parts.append(f"<index>{m['index']}</index>")
                parts.append(f"<antenna-mask>{m['b64']}</antenna-mask>")
                parts.append("</mplane-supported-trx-control-masks>")
            parts.append("</mplane-trx-control-txarr-capability-info>")
            parts.append("</tx-arrays>")
        # Current config
        for aname, cfg in self.current_config.items():
            parts.append("<tx-array-antenna-mask-config>")
            parts.append(f"<array-name>{aname}</array-name>")
            parts.append(f"<antenna-bitmask-index>{cfg['index']}</antenna-bitmask-index>")
            parts.append("</tx-array-antenna-mask-config>")
        parts.append("</user-plane-configuration>")
        return "".join(parts)


# ---------------------------------------------------------------------------
# gNB internal antenna control client (localhost:9999)
# ---------------------------------------------------------------------------

class GnbCtrlClient:
    """Connects to the gNB internal test socket to apply antenna masks."""

    def __init__(self, host, port):
        self.host = host
        self.port = port
        self.sock = None

    def connect(self):
        for attempt in range(30):
            try:
                self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                self.sock.settimeout(5.0)
                self.sock.connect((self.host, self.port))
                self._recv()  # consume welcome
                log.info("Connected to gNB control at %s:%d", self.host, self.port)
                return True
            except (socket.error, ConnectionRefusedError) as e:
                log.debug("gNB ctrl connect attempt %d/30: %s", attempt + 1, e)
                if self.sock:
                    self.sock.close()
                    self.sock = None
                time.sleep(2)
        log.warning("Cannot connect to gNB control at %s:%d", self.host, self.port)
        return False

    def send_command(self, cmd):
        if not self.sock:
            if not self.connect():
                return "ERROR: not connected"
        try:
            self.sock.sendall((cmd + "\n").encode())
            time.sleep(0.1)
            return self._recv()
        except (socket.error, BrokenPipeError):
            log.warning("gNB ctrl connection lost, reconnecting...")
            self.sock = None
            if not self.connect():
                return "ERROR: reconnect failed"
            self.sock.sendall((cmd + "\n").encode())
            time.sleep(0.1)
            return self._recv()

    def _recv(self):
        self.sock.settimeout(0.5)
        data = b""
        try:
            while True:
                chunk = self.sock.recv(4096)
                if not chunk:
                    break
                data += chunk
        except socket.timeout:
            pass
        return data.decode("utf-8", errors="ignore")

    def close(self):
        if self.sock:
            try:
                self.sock.sendall(b"QUIT\n")
            except Exception:
                pass
            self.sock.close()
            self.sock = None


# ---------------------------------------------------------------------------
# NETCONF Session Handler
# ---------------------------------------------------------------------------

class NetconfSession:
    """Handle a single NETCONF session on an SSH channel."""

    def __init__(self, channel, capability, gnb_ctrl):
        self.channel = channel
        self.cap = capability
        self.gnb = gnb_ctrl
        self.session_id = 1
        self.running = True

    def run(self):
        try:
            self._send_hello()
            hello_msg = self._recv_message()
            if hello_msg is None:
                log.warning("Client did not send hello")
                return
            log.info("NETCONF hello exchange complete")

            while self.running:
                msg = self._recv_message()
                if msg is None:
                    break
                self._process_rpc(msg)
        except Exception as e:
            log.exception("Session error: %s", e)
        finally:
            try:
                self.channel.close()
            except Exception:
                pass
            log.info("NETCONF session closed")

    def _send_hello(self):
        hello = (
            '<?xml version="1.0" encoding="UTF-8"?>'
            f'<hello xmlns="{NS_NC}">'
            "<capabilities>"
            "<capability>urn:ietf:params:netconf:base:1.0</capability>"
            "<capability>urn:ietf:params:netconf:capability:writable-running:1.0</capability>"
            "</capabilities>"
            f"<session-id>{self.session_id}</session-id>"
            "</hello>"
        )
        self._send_message(hello)

    def _send_message(self, xml_str):
        data = (xml_str + NETCONF_END).encode("utf-8")
        self.channel.sendall(data)

    def _recv_message(self):
        buf = b""
        while self.running:
            try:
                self.channel.settimeout(300)
                chunk = self.channel.recv(65536)
                if not chunk:
                    return None
                buf += chunk
                text = buf.decode("utf-8", errors="ignore")
                if NETCONF_END in text:
                    return text.split(NETCONF_END)[0].strip()
            except socket.timeout:
                continue
            except Exception:
                return None
        return None

    def _process_rpc(self, xml_str):
        # Strip XML declaration if present
        xml_str = re.sub(r"<\?xml[^?]*\?>", "", xml_str).strip()
        try:
            root = ET.fromstring(xml_str)
        except ET.ParseError as e:
            log.warning("Malformed RPC: %s", e)
            self._send_error("0", "malformed-message", "rpc")
            return

        msg_id = root.get("message-id", "0")

        # Find the operation element
        for child in root:
            local_tag = child.tag.split("}")[-1] if "}" in child.tag else child.tag

            if local_tag == "edit-config":
                self._handle_edit_config(msg_id, child)
            elif local_tag == "get":
                self._handle_get(msg_id, child)
            elif local_tag == "get-config":
                self._handle_get_config(msg_id, child)
            elif local_tag == "close-session":
                self._send_ok(msg_id)
                self.running = False
            elif local_tag == "lock" or local_tag == "unlock":
                self._send_ok(msg_id)
            else:
                self._send_error(msg_id, "operation-not-supported", "protocol")
            return  # Only process first operation

    def _handle_edit_config(self, msg_id, edit_elem):
        """Handle edit-config for tx-array-antenna-mask-config."""
        # Find <config> element
        config_elem = None
        for child in edit_elem:
            local = child.tag.split("}")[-1] if "}" in child.tag else child.tag
            if local == "config":
                config_elem = child
                break

        if config_elem is None:
            self._send_error(msg_id, "missing-element", "protocol",
                             "Missing <config> element")
            return

        # Find tx-array-antenna-mask-config under user-plane-configuration
        mask_config = None
        for elem in config_elem.iter():
            local = elem.tag.split("}")[-1] if "}" in elem.tag else elem.tag
            if local == "tx-array-antenna-mask-config":
                mask_config = elem
                break

        if mask_config is None:
            # No antenna mask config in this edit - accept silently
            self._send_ok(msg_id)
            return

        # Extract array-name, antenna-bitmask-index or antenna-bitmask
        array_name = None
        bitmask_index = None
        bitmask_b64 = None

        for child in mask_config:
            local = child.tag.split("}")[-1] if "}" in child.tag else child.tag
            if local == "array-name":
                array_name = child.text.strip() if child.text else None
            elif local == "antenna-bitmask-index":
                bitmask_index = int(child.text.strip()) if child.text else None
            elif local == "antenna-bitmask":
                bitmask_b64 = child.text.strip() if child.text else None

        if not array_name:
            self._send_error(msg_id, "missing-element", "application",
                             "Missing array-name in tx-array-antenna-mask-config")
            return

        # Resolve the mask
        resolved = None
        if bitmask_index is not None:
            resolved, err = self.cap.resolve_index(array_name, bitmask_index)
            if not resolved:
                self._send_error(msg_id, "invalid-value", "application", err)
                return
            log.info("M-Plane edit-config: array=%s, antenna-bitmask-index=%d",
                     array_name, bitmask_index)
        elif bitmask_b64 is not None:
            mask_int = b64_to_mask(bitmask_b64)
            resolved, err = self.cap.resolve_mask(array_name, mask_int)
            if not resolved:
                self._send_error(msg_id, "invalid-value", "application", err)
                return
            log.info("M-Plane edit-config: array=%s, antenna-bitmask=%s (0x%02X)",
                     array_name, bitmask_b64, mask_int)
        else:
            self._send_error(msg_id, "missing-element", "application",
                             "Need antenna-bitmask or antenna-bitmask-index")
            return

        # Apply to gNB via internal control socket and verify it succeeded.
        # The gNB returns "OK: TX mask 0xAA -> 0xBB" on success; the socket
        # wrapper returns "ERROR: ..." on transport failure and the gNB itself
        # returns "ERROR: ..." for invalid input.  Returning <ok/> when the
        # underlying control failed would let the M-Plane state and the gNB
        # state diverge.
        mask_int = resolved["int"]
        resp = self.gnb.send_command(f"SET TX MASK 0x{mask_int:02X}")
        resp_clean = (resp or "").strip()

        if not resp_clean.upper().startswith("OK"):
            log.warning("gNB rejected mask 0x%02X: %r", mask_int, resp_clean)
            self._send_error(
                msg_id,
                "operation-failed",
                "application",
                f"gNB control did not confirm mask change: {resp_clean or 'no response'}",
            )
            return

        log.info("Applied mask 0x%02X to gNB: %s", mask_int, resp_clean)

        # Update current config only after the gNB confirmed.
        self.cap.current_config[array_name] = {
            "index": resolved["index"],
            "mask_int": mask_int,
        }
        # Note: the production O1 path used by this paper runs through
        # docker/o1-adapter/o1_antenna_bridge.py; this legacy gNB-side
        # M-Plane server keeps the same OK-required check for parity.

        # Log WG4 mplane-trx-control-ant-mask-update notification
        log.info("WG4 mplane-trx-control-ant-mask-update: "
                 "array-name='%s', array-type=TX-ARRAY, "
                 "updated-antenna-mask-index=%d, updated-antenna-mask=%s",
                 array_name, resolved["index"], resolved["b64"])

        self._send_ok(msg_id)

    def _handle_get(self, msg_id, get_elem):
        """Handle <get> - return operational data (capability + current config)."""
        data_xml = self.cap.get_capability_xml()
        reply = (
            f'<rpc-reply xmlns="{NS_NC}" message-id="{msg_id}">'
            f"<data>{data_xml}</data>"
            "</rpc-reply>"
        )
        self._send_message(reply)

    def _handle_get_config(self, msg_id, get_elem):
        """Handle <get-config> - return running config (current mask config)."""
        parts = [f'<user-plane-configuration xmlns="{NS_UPLANE}">']
        for aname, cfg in self.cap.current_config.items():
            parts.append("<tx-array-antenna-mask-config>")
            parts.append(f"<array-name>{aname}</array-name>")
            parts.append(f"<antenna-bitmask-index>{cfg['index']}</antenna-bitmask-index>")
            parts.append("</tx-array-antenna-mask-config>")
        parts.append("</user-plane-configuration>")
        reply = (
            f'<rpc-reply xmlns="{NS_NC}" message-id="{msg_id}">'
            f'<data>{"".join(parts)}</data>'
            "</rpc-reply>"
        )
        self._send_message(reply)

    def _send_ok(self, msg_id):
        reply = (
            f'<rpc-reply xmlns="{NS_NC}" message-id="{msg_id}">'
            "<ok/>"
            "</rpc-reply>"
        )
        self._send_message(reply)

    def _send_error(self, msg_id, error_tag, error_type, error_message=""):
        msg_part = ""
        if error_message:
            msg_part = f"<error-message>{error_message}</error-message>"
        reply = (
            f'<rpc-reply xmlns="{NS_NC}" message-id="{msg_id}">'
            "<rpc-error>"
            f"<error-type>{error_type}</error-type>"
            f"<error-tag>{error_tag}</error-tag>"
            "<error-severity>error</error-severity>"
            f"{msg_part}"
            "</rpc-error>"
            "</rpc-reply>"
        )
        self._send_message(reply)


# ---------------------------------------------------------------------------
# SSH Server Interface
# ---------------------------------------------------------------------------

class MplaneSSHServer(paramiko.ServerInterface):
    def __init__(self):
        self.subsystem_event = threading.Event()

    def check_channel_request(self, kind, chanid):
        if kind == "session":
            return paramiko.OPEN_SUCCEEDED
        return paramiko.OPEN_FAILED_ADMINISTRATIVELY_PROHIBITED

    def get_allowed_auths(self, username):
        return "password"

    def check_auth_password(self, username, password):
        if username == MPLANE_USER and password == MPLANE_PASS:
            return paramiko.AUTH_SUCCESSFUL
        return paramiko.AUTH_FAILED

    def check_channel_subsystem_request(self, channel, name):
        if name == "netconf":
            self.subsystem_event.set()
            return True
        return False


# ---------------------------------------------------------------------------
# Main Server
# ---------------------------------------------------------------------------

def handle_connection(client_sock, addr, host_key, capability, gnb_ctrl):
    """Handle a single SSH/NETCONF client connection."""
    log.info("SSH connection from %s:%d", addr[0], addr[1])
    transport = None
    try:
        transport = paramiko.Transport(client_sock)
        transport.add_server_key(host_key)
        server = MplaneSSHServer()
        transport.start_server(server=server)

        channel = transport.accept(30)
        if channel is None:
            log.warning("No channel opened by client")
            return

        # Wait for NETCONF subsystem request
        if not server.subsystem_event.wait(10):
            log.warning("Client did not request netconf subsystem")
            channel.close()
            return

        session = NetconfSession(channel, capability, gnb_ctrl)
        session.run()

    except paramiko.SSHException as e:
        log.warning("SSH error from %s: %s", addr[0], e)
    except Exception as e:
        log.exception("Connection error from %s: %s", addr[0], e)
    finally:
        if transport:
            try:
                transport.close()
            except Exception:
                pass


def main():
    log.info("O-RU M-Plane NETCONF server starting")
    log.info("  Port: %d", MPLANE_PORT)
    log.info("  gNB control: %s:%d (internal)", GNB_CTRL_HOST, GNB_CTRL_PORT)

    # Load/generate SSH host key
    if os.path.exists(SSH_HOST_KEY):
        host_key = paramiko.RSAKey.from_private_key_file(SSH_HOST_KEY)
        log.info("Loaded SSH host key from %s", SSH_HOST_KEY)
    else:
        host_key = paramiko.RSAKey.generate(2048)
        os.makedirs(os.path.dirname(SSH_HOST_KEY), exist_ok=True)
        host_key.write_private_key_file(SSH_HOST_KEY)
        log.info("Generated SSH host key at %s", SSH_HOST_KEY)

    # Load WG4 capability
    capability = Capability()
    capability.load(WG4_CAPABILITY_FILE)
    for name, arr in capability.tx_arrays.items():
        log.info("  tx-array '%s': %d supported masks", name, len(arr["masks"]))
        for m in arr["masks"]:
            log.info("    index=%d: %s (0x%02X, %d active)",
                     m["index"], m["b64"], m["int"], bin(m["int"]).count("1"))

    # Connect to gNB internal control
    gnb_ctrl = GnbCtrlClient(GNB_CTRL_HOST, GNB_CTRL_PORT)
    gnb_ctrl.connect()

    # Create server socket
    server_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server_sock.bind(("0.0.0.0", MPLANE_PORT))
    server_sock.listen(5)
    server_sock.settimeout(1.0)

    shutdown = threading.Event()

    def handle_signal(signum, frame):
        log.info("Shutdown signal received")
        shutdown.set()

    signal.signal(signal.SIGTERM, handle_signal)
    signal.signal(signal.SIGINT, handle_signal)

    log.info("M-Plane NETCONF server ready on port %d", MPLANE_PORT)

    while not shutdown.is_set():
        try:
            client_sock, addr = server_sock.accept()
            t = threading.Thread(
                target=handle_connection,
                args=(client_sock, addr, host_key, capability, gnb_ctrl),
                daemon=True,
            )
            t.start()
        except socket.timeout:
            continue
        except Exception as e:
            if not shutdown.is_set():
                log.exception("Accept error: %s", e)

    log.info("Shutting down...")
    server_sock.close()
    gnb_ctrl.close()


if __name__ == "__main__":
    main()