from .ssh import SSHClient
import os

class RouterManager:
    def __init__(self, config):
        self.host = config.get('router.host')
        self.user = config.get('router.user', 'root')
        self.ssh_key = config.get('router.ssh_key')
        self.password = os.environ.get("ROUTER_PASS")
        self.client = SSHClient(self.host, self.user, self.ssh_key, self.password)

    def run_uci(self, cmds):
        res = self.client.run(cmds)
        if res.returncode != 0:
            raise RuntimeError(f"UCI command failed: {res.stderr}")
        return res

    def add_forward(self, name, proto, ext_port, int_ip, int_port):
        config_name = name.lower().replace(" ", "_").replace("-", "_")
        
        cmds = f"uci set firewall.{config_name}=redirect; " \
               f"uci set firewall.{config_name}.name='{name}'; " \
               f"uci set firewall.{config_name}.proto='{proto}'; " \
               f"uci set firewall.{config_name}.src='wan'; " \
               f"uci set firewall.{config_name}.dest='lan'; " \
               f"uci set firewall.{config_name}.src_dport='{ext_port}'; " \
               f"uci set firewall.{config_name}.dest_ip='{int_ip}'; " \
               f"uci set firewall.{config_name}.dest_port='{int_port}'; " \
               f"uci set firewall.{config_name}.target='DNAT'; " \
               f"uci commit firewall; " \
               f"/etc/init.d/firewall reload"
        
        self.run_uci(cmds)

    def remove_forward(self, section_name):
        cmds = f"uci delete firewall.{section_name}; uci commit firewall; /etc/init.d/firewall reload"
        self.run_uci(cmds)

    def list(self):
        # Get sections
        # Use -n to prevent SSH from consuming stdin if we were in a loop (here we use split)
        res = self.client.run("uci show firewall | grep '=redirect' | cut -d. -f2 | cut -d= -f1 | sort | uniq")
        sections = res.stdout.strip().split('\n')
        
        results = []
        for section in sections:
            if not section: continue
            name = self.client.run(f"uci get firewall.{section}.name").stdout.strip()
            proto = self.client.run(f"uci get firewall.{section}.proto").stdout.strip()
            port = self.client.run(f"uci get firewall.{section}.src_dport").stdout.strip()
            dest = self.client.run(f"uci get firewall.{section}.dest_ip").stdout.strip()
            results.append({
                "section": section,
                "name": name,
                "proto": proto,
                "port": port,
                "dest": dest
            })
        return results