#!/usr/bin/env python3
# Copyright 2020 Canonical Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
#     Unless required by applicable law or agreed to in writing, software
#     distributed under the License is distributed on an "AS IS" BASIS,
#     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#     See the License for the specific language governing permissions and
#     limitations under the License.
import sys

sys.path.append("lib")
from ops.charm import CharmBase
from ops.framework import StoredState
from ops.main import main
from ops.model import (
    ActiveStatus,
    BlockedStatus,
    MaintenanceStatus,
    WaitingStatus,
    ModelError,
)
import subprocess
import logging
import asyncio
import random

logger = logging.getLogger(__name__)


class EnodebCharm(CharmBase):
    state = StoredState()

    def __init__(self, *args):
        super().__init__(*args)
        # An example of setting charm state
        # that's persistent across events
        self.state.set_default(is_started=False)
        if not self.state.is_started:
            self.state.is_started = True

        # Register all of the events we want to observe
        self.framework.observe(self.on.config_changed, self.on_config_changed)
        self.framework.observe(self.on.install, self.on_install)
        self.framework.observe(self.on.start, self.on_start)
        self.framework.observe(self.on.upgrade_charm, self.on_upgrade_charm)

        self.framework.observe(self.on.attach_ue_action, self.on_attach_ue_action)
        self.framework.observe(self.on.detach_ue_action, self.on_detach_ue_action)
        self.framework.observe(self.on.remove_default_gw_action, self.on_remove_default_gw_action)

        self.framework.observe(self.on.agw_relation_changed, self.agw_relation_changed)

    def on_config_changed(self, event):
        """Handle changes in configuration"""

    def on_start(self, event):
        """Called when the charm starts"""
        self.unit.status = ActiveStatus()

    def on_install(self, event):
        """Called when the charm is being installed"""

    def on_upgrade_charm(self, event):
        """Upgrade the charm."""

    def agw_relation_changed(self, event):
        """Register to AGW (EPC)."""
        self.unit.status = MaintenanceStatus("Getting MagmaGW data from relation")
        relation = self.model.get_relation("agw")
        if relation is None:
            event.defer()
            return
        self._unregister()
        mme_addr = relation.data[event.unit].get("mme-addr", None)
        status = "Not registered"
        if mme_addr is not None:
            self.unit.status = MaintenanceStatus(f"Relation data {dict(relation.data[event.unit])}")
            gtp_bind_addr = self._get_bind_address()
            s1c_bind_addr = gtp_bind_addr
            command = " ".join(
                [
                    "/home/ubuntu/srsLTE/build/srsenb/src/srsenb",
                    "--enb.name=dummyENB01",
                    "--enb.mcc=901",
                    "--enb.mnc=70",
                    "--enb.mme_addr={}".format(mme_addr),
                    "--enb.gtp_bind_addr={}".format(gtp_bind_addr),
                    "--enb.s1c_bind_addr={}".format(s1c_bind_addr),
                    "--enb_files.rr_config=/configzmq/rr.conf",
                    "--enb_files.sib_config=/configzmq/sib.conf",
                    "--enb_files.drb_config=/configzmq/drb.conf",
                    "/configzmq/enb.conf",
                    "--rf.device_name=zmq",
                    "--rf.device_args='fail_on_disconnect=true,tx_port=tcp://*:2000,rx_port=tcp://localhost:2001,id=enb,base_srate=23.04e6'",
                ]
            )
            self._run_daemon(command)
            status = "Registered"
        self.unit.status = ActiveStatus(status)

    def _get_bind_address(self):
        output = subprocess.run(["hostname", "-I"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        return output.stdout.decode("utf-8").split(" ")[1]

    def _unregister(self):
        """Unregister"""
        self.unit.status = MaintenanceStatus("Unregistering")
        subprocess.run("sudo killall -s KILL srsenb", shell=True)

    def on_attach_ue_action(self, event):
        """Attach User Emulator to EnodeB."""
        try:
            usim_imsi = event.params["usim-imsi"]
            usim_k = event.params["usim-k"]
            usim_opc = event.params["usim-opc"]
            command = " ".join(
                [
                    "sudo",
                    "/home/ubuntu/srsLTE/build/srsue/src/srsue",
                    "--usim.imsi={}".format(usim_imsi),
                    "--usim.k={}".format(usim_k),
                    "--usim.algo=milenage",
                    "--usim.opc={}".format(usim_opc),
                    "--nas.apn=oai.ipv4",
                    "--rf.device_name=zmq",
                    "--rf.device_args='tx_port=tcp://*:2001,rx_port=tcp://localhost:2000,id=ue,base_srate=23.04e6'",
                    "/configzmq/ue.conf",
                ]
            )
            process = self._run_daemon(command)
            event.set_results(
                {"status": "ok", "pid": process.pid}
            )
        except subprocess.CalledProcessError as ex:
            event.fail(ex)

    def on_detach_ue_action(self, event):
        """Detach UE action"""
        try:
            command = "sudo killall -s KILL srsue"
            output = subprocess.check_output(command, shell=True)
            event.set_results(
                {"status": "ok", "message": "Detached successfully", "output": output}
            )
        except subprocess.CalledProcessError as e:
            event.fail("Command error: {}".format(e.output))
        except Exception as e:
            event.fail(e)

    def on_remove_default_gw_action(self, event):
        """Remove default gw"""
        try:
            command = "sudo route del default"
            output = subprocess.check_output(command, shell=True)
            event.set_results(
                {"status": "ok", "message": "Default route removed!", "output": output}
            )
        except subprocess.CalledProcessError as e:
            event.fail("Command error: {}".format(e.output))
        except Exception as e:
            event.fail(e)

    def _run_daemon(self, cmd):
        log_basename = random.randint(1000, 100000)
        stdout_file = f"/tmp/{log_basename}.stdout"
        stderr_file = f"/tmp/{log_basename}.stderr"
        with open(stdout_file, "wb") as o:
            with open(stderr_file, "wb") as e:
                return subprocess.Popen(cmd, shell=True, stdout=o, stderr=e)


if __name__ == "__main__":
    main(EnodebCharm)
