import os
from openvpn import Openvpn
import asyncio
from tornado.template import Template
import signal


class OManager:
    def __init__(self, base_folder, base_port=8001, loop=None):
        self.base_folder = base_folder
        self.base_port = base_port
        self.ops = []
        self.idx = 0
        self.running_idx = []
        self.run_task = []
        self.PID = None 
        if loop:
            self.loop = loop
        else:
            self.loop = asyncio.get_event_loop()

    def new_op(self, cfg_fp, name=None, additional_cfg={}):
        folder_path = os.path.join(self.base_folder, f"session{self.idx}")
        if not os.path.isdir(folder_path):
            os.makedirs(folder_path)
        if not name:
            name = f"openvpn-{self.idx}"
        os.system(f"groupadd vpn{self.idx}")
        op = Openvpn(cfg_fp, self.idx, folder_path,
                     f"{self.base_port + self.idx}", "script.sh.template", "3proxy.cfg.template", name=name,
                     additional_cfg=additional_cfg)
        self.ops.append(op)
        self.idx += 1
        return op

    def generate_lb_cfg(self):
        lb_cfg_fp = os.path.join(
            self.base_folder, "go-socks-lb.yml")
        with open("go-socks-lb.yml.template", "r") as template_f:
            buf = template_f.read()
        template = Template(buf)
        items = {
            "running_idx": self.running_idx,
            "port": f"{1080 + self.idx}"
        }
        lb_cfg = template.generate(**items)
        with open(lb_cfg_fp, "wb") as lb_cfg_f:
            lb_cfg_f.write(lb_cfg)
        return lb_cfg_fp

    def start_op(self, idx):
        self.ops[idx].start()
        if idx not in self.running_idx:
            self.running_idx.append(idx)
        self.reset_lb()

    def stop_op(self, idx):
        self.loop.create_task(self.ops[idx].stop())
        if idx in self.running_idx:
            self.running_idx.remove(idx)
        self.reset_lb()

    def get_all_ops(self):
        return self.ops

    def start_all(self):
        for op in self.ops:
            op.start()
            idx = op.idx
            if idx not in self.running_idx:
                self.running_idx.append(idx)
        self.reset_lb()

    def stop_all(self):
        loop = self.loop
        for op in self.ops:
            loop.create_task(op.stop())
        self.running_idx = []
        self.reset_lb()

    def reset_lb(self):
        try:
            if self.PID:
                os.kill(self.PID, signal.SIGINT)
        except Exception as err:
            print("kill failed:", err)
        for task in self.run_task:
            task.cancel()
        self.run_task = []
        lb_cfg_fp = self.generate_lb_cfg()
        lb_cmd = ["go-socks-lb/go-socks-lb", "-config",
                           lb_cfg_fp, "-bind", "0.0.0.0:7000"]
        print("lb_cmd", lb_cmd)
        self.run_task.append(self.loop.create_task(
            self.run_cmd(lb_cmd)))

    async def run_cmd(self, cmd):
        while True:
            print("Manager trying to start go")
            proc = await asyncio.create_subprocess_exec(
                    cmd[0], *cmd[1:], 
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE)
            print("Manager started", cmd)
            self.proc = proc
            self.PID = proc.pid
            print(f"pid: {self.PID}")
            stdout, stderr = await proc.communicate()
            print(f'[{cmd!r} exited with {proc.returncode}]')
            if stdout:
                print(f'[stdout]\n{stdout.decode()}')
            if stderr:
                print(f'[stderr]\n{stderr.decode()}')
            await asyncio.sleep(5)


if __name__ == "__main__":
    folder_fp = "/home/mantao/Desktop/t/"
    cfg_fp = "/home/mantao/Desktop/t/TCP_Files/UK2-TCP.ovpn"
    om = OManager(folder_fp)
    om.new_op(cfg_fp, "op1", {
        "auth-user-pass": "/home/mantao/Desktop/t/fast.txt"})