import fcntl
import os
import re
import socket
import struct
import subprocess
import sys
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from typing import Any, TypeAlias

TUNNEL_ENABLE_FMT_PARTS = [
    "ssh -fN {direction} *:{port_from}:*:{port_to} {host}",
    "-o ControlMaster=yes",
    "-o ControlPath={control_path}",
    "-o ServerAliveInterval=60",
    "-o ExitOnForwardFailure=yes",
]
TUNNEL_ENABLE_FMT = " ".join(TUNNEL_ENABLE_FMT_PARTS)
TUNNEL_DISABLE_FMT = "ssh -O exit -S {control_path} {host}"
SUGAR_EXAMPLES = (
    (":1111 -> example:1111", ":1111 -> example:1111"),
    ("1111 -> example      ", ":1111 -> example:1111"),
    ("5555 example:1111    ", ":5555 -> example:1111"),
    ("example:1111         ", ":1111 -> example:1111"),
    ("example 1111         ", ":1111 -> example:1111"),
    ("1111 example         ", ":1111 -> example:1111"),
    (":1111 <- example:1111", ":1111 <- example:1111 [REVERSE]"),
    ("1111 <- example      ", ":1111 <- example:1111 [REVERSE]"),
    ("5555 example:1111 --r", ":5555 <- example:1111 [REVERSE]"),
    ("example:1111 reversed", ":1111 <- example:1111 [REVERSE]"),
    ("1111 example --reverse", ":1111 <- example:1111 [REVERSE]"),
)
HELP_COMMANDS = """This is tool to manage tunnels.
Allowed example commands:
- 1111 -> example:999
"""
HELP_COMMANDS += "\n".join(f"- {k} :: sugar for {val}" for k, val in SUGAR_EXAMPLES)
SOCKET_FNAME_PATTERN = r"ssh-tunnel--(\d+)--([^-]+)-(\d+)--(L|R).sock"
SOCKET_FNAME_FMT = r"ssh-tunnel--{local_port}--{host}-{port}-{direction}.sock"
DRY = bool(int(os.environ.get("dry", "0")))
D = bool(int(os.environ.get("d", "0")))
Printer: TypeAlias = Callable[[str], Any]
Either = tuple
Ignore = lambda arg: None  # noqa
debug = (lambda arg: print("#", arg)) if D else Ignore


## system
def _exec_command(cmd) -> tuple[Exception | None, str | None]:
    try:
        result = subprocess.run(
            cmd,
            shell=True,
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )
        return None, (result.stdout + "\n" + (result.stderr or "")).strip()
    except subprocess.CalledProcessError as ex:
        return ex, None


def get_ip2(debug: Printer = Ignore) -> tuple[Exception | None, str | None]:
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(("8.8.8.8", 80))
        res = s.getsockname()[0]
        s.close()
        return None, res
    except Exception as e:
        debug(f"Getting IP: failed: {e}")
        return Exception(f"Failed to get IP: {e}"), None


def get_ip(interface="eth0", *, debug: Printer = Ignore) -> tuple[Exception | None, str | None]:
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        interface_enc = interface[:15].encode()
        file_descriptor = sock.fileno()
        request = 0x8915  # SIOCGIFADDR
        arg = struct.pack("256s", interface_enc)
        inet = fcntl.ioctl(file_descriptor, request, arg)
        inet_part = inet[20:24]
        ip = socket.inet_ntoa(inet_part)
        return None, ip
    except OSError as e:
        debug(f"Getting IP for {interface}: failed: {e}")
        return get_ip2(debug)


def get_hostname_pretty() -> str | None:
    _, name = _exec_command("hostnamectl --pretty")
    return name


def get_hostname(use_pretty_if_exist: bool = True) -> str | None:
    if use_pretty_if_exist and (name := get_hostname_pretty()):
        return name

    err, name = _exec_command("hostname")
    return name


def get_host_info(debug: Printer = Ignore):
    hostname = get_hostname()
    err, ip = get_ip(debug=debug)
    ip = ip or "??"
    return f"[host {hostname} ({ip})]"


## tunnels
@dataclass
class TunnelInfo:
    local_port: int
    host: str
    port: int
    forward: bool

    def prettify(self):
        if self.forward:
            res = f":{self.local_port} -> {self.host}:{self.port}"
        else:
            res = f":{self.local_port} <- {self.host}:{self.port} [REVERSE]"
        return res

    def to_dict(self):
        port_from, port_to = (self.local_port, self.port) if self.forward else (self.port, self.local_port)
        return {
            "local_port": self.local_port,
            "host": self.host,
            "port": self.port,
            "direction": "-L" if self.forward else "-R",
            "port_from": port_from,
            "port_to": port_to,
        }

    def make_socket_fname(self):
        return SOCKET_FNAME_FMT.format(**self.to_dict())

    def t(self):
        return (self.local_port, self.host, self.port)


def make_control_path(tun: TunnelInfo):
    return Path.home() / ".ssh" / tun.make_socket_fname()


def make_tunnel_command(tun: TunnelInfo, enable: bool):
    cmd_fmt = TUNNEL_ENABLE_FMT if enable else TUNNEL_DISABLE_FMT
    tun_cmd_args = tun.to_dict()
    tun_cmd_args["control_path"] = str(make_control_path(tun))
    cmd = cmd_fmt.format(**tun_cmd_args)
    return cmd


def parse_tunnel(sock: Path) -> TunnelInfo | None:
    mt = re.match(SOCKET_FNAME_PATTERN, sock.name)
    if not mt:
        return None
    local_port, host, port, direction = mt.group(1), mt.group(2), mt.group(3), mt.group(4)
    forward = direction == "L"
    res = TunnelInfo(int(local_port), host, int(port), forward)
    debug(f"Parsed tunnel: {res}")
    return res


def get_active_tunnels() -> list[TunnelInfo]:
    socks = (Path.home() / ".ssh").glob("*.sock")
    tunnels = list(filter(None, map(parse_tunnel, socks)))
    return tunnels


REVERSE_TOKENS = [
    "--reversed",
    "reversed",
    "--reverse",
    "reverse",
    "--r",
    "--R",
    "<-",
    "-r",
    "-r",
]


def remove_all(text, tokens):
    for token in tokens:
        text = text.replace(token, "").strip()
    return text


def parse_tunnel_args(cmd_parts: list[str], forward=True) -> TunnelInfo:
    assert isinstance(cmd_parts, list), cmd_parts
    debug(f"parsing tunnel args: {cmd_parts}")

    if len(cmd_parts) == 1:
        host_port = cmd_parts[0]
        host, port = host_port.split(":", 1)

        local_port, host = port, host_port
    else:
        local_port, host = cmd_parts

    if host.isnumeric():
        host, local_port = local_port, host

    if ":" in host:
        host, port = host.split(":", 1)
    else:
        port = local_port
    host = host.strip()
    local_port = local_port.lstrip(":")
    res = TunnelInfo(int(local_port), host, int(port), forward=forward)
    debug(f"Parsed tunnel: {res}")
    return res


def parse_tunnel_command(cmd: str) -> TunnelInfo:
    assert isinstance(cmd, str), cmd
    debug(f"parsing cmd: {cmd}")
    cmd = cmd.strip()

    cmd_fix = remove_all(cmd, REVERSE_TOKENS)
    forward = cmd_fix == cmd
    cmd = cmd_fix
    cmd = cmd.replace("->", " ")
    debug(f"After forward preprocessing cmd: {cmd}")
    parts = list(filter(None, re.split(r"\s+", cmd)))
    return parse_tunnel_args(parts, forward)


def test_parse_tunnel_command():
    for cmd, expected in SUGAR_EXAMPLES:
        cmd = cmd.strip()
        actual = parse_tunnel_command(cmd).prettify()
        if actual == expected:
            print(f"OK: `{cmd}` -> `{actual}`")
            continue
        print(f"cmd:      `{cmd}`")
        print(f"actual:   `{actual}`")
        print(f"expected: `{expected}`")
        raise ValueError("Not equal")


class Tunneler:
    def __init__(self, debug: Printer, dry: bool):
        self.debug = debug
        self.dry = dry

    def is_active(self, tun: TunnelInfo) -> bool:
        control_path = make_control_path(tun)
        return bool(control_path.exists())

    def prettify(self, tun: TunnelInfo, extra: str | None = None) -> str:
        status = "ACTIVE" if self.is_active(tun) else "DISABLED"
        extra = f", {extra}" if extra else ""
        return f"tunnel[port={tun.port}, status={status}{extra}]"

    def toggle(self, tun: TunnelInfo, *, enable: bool):
        is_active = self.is_active(tun)
        if is_active and enable:
            print("Tunnel is already enabled!")
            return
        if not is_active and not enable:
            print("Tunnel is already not active!")
            return
        if enable and tun.host is None:
            print("No `donor_host` installed, can't forward port!")
            return
        w_toggle = "enabling" if enable else "disabling"
        msg_pref = f"Tunnel '{tun.prettify()}': {w_toggle}"
        debug(f"{msg_pref}: going to start")
        # todo fix: forbid to create tunnel which is already exist
        sys_cmd = make_tunnel_command(tun, enable=enable)
        if self.dry:
            print(f"DRY: exit (skipping to execute command: '{sys_cmd}')")
            return
        debug(f"{msg_pref}: execution: {sys_cmd}")
        ex, text = _exec_command(sys_cmd)
        if not ex:
            print(f"{msg_pref}: done")
            return
        err = ex.args[0]
        print(f"{msg_pref}: FAILED")
        if err == 255:
            err = "255 (probably it means 'Connection refused')"
        print(err)

    def get_active_tunnels(self) -> list[TunnelInfo]:
        return get_active_tunnels()


def _parse_port(args):
    if not args or len(args) != 1:
        return None
    arg = args[0]
    if isinstance(arg, int):
        return arg
    return int(arg) if arg.isnumeric() else None


def _kwargs_as_args(kwargs):
    return [f"--{kw}" for kw, kw_val in kwargs.items() if kw_val]


class TunnelerCLI:
    def __init__(self):
        self.tunneler = Tunneler(debug, DRY)

    def tunnels(self, args):
        active_tunnels = self.tunneler.get_active_tunnels()
        host_info = get_host_info(self.tunneler.debug)
        if not active_tunnels:
            return f"No active tunnels {host_info}"
        res_lines = [
            f"Active tunnels {host_info}:",
            # todo show target
            *[f"- {tun.prettify()}" for tun in active_tunnels],
        ]
        res = "\n".join(res_lines)
        return res

    def _toggle(self, enable: bool, args):
        tun = parse_tunnel_command(" ".join(map(str, args)))
        self.tunneler.toggle(tun, enable=enable)

    def parse(self, args):
        assert len(args) == 1
        cmd = args[0]
        parsed_cmd = parse_tunnel_command(cmd)
        return f"your command: '{cmd}', parsed: {parsed_cmd}"

    def down(self, args):
        tuns = get_active_tunnels()
        if not tuns:
            return "No tunnels opened"
        port = _parse_port(args)

        if len(tuns) == 1 and not args:
            tun = tuns[0]
            print(f"Selected the one opened tunnel: {tun.prettify()}")
            return self.tunneler.toggle(tun, enable=False)

        if not args:
            return "No tunnel specified"

        if port:
            tuns_port = [tc for tc in tuns if tc.local_port == port]
            if not tuns_port:
                return f"Port {port}: tunnel is not opened"
            assert len(tuns_port) == 1
            return self.tunneler.toggle(tuns_port[0], enable=False)

        tun = parse_tunnel_command(" ".join(map(str, args)))
        if tun not in tuns:
            return f"Tunnel {tun} is not opened"

        return self.tunneler.toggle(tun, enable=False)

    def up(self, args):
        return self._toggle(True, args)

    # fmt: off
    def enable(self, args): return self.up(*args)
    def disable(self, args): return self.down(*args)
    def test(self, args): test_parse_tunnel_command()
    def tuns(self, args): return self.tunnels([])
    def help(self, args): return HELP_COMMANDS
    def __str__(self): return self.tunnels([])
    # fmt: on

    def invoke(self, args):
        if len(args) == 0:
            res = str(self)
        else:
            arg = args[0]
            if arg.isnumeric():
                cmd = "up"
            else:
                cmd, *args = args

            commands = [field for field in dir(self) if not field.startswith("_") and field != "invoke"]
            if cmd not in commands:
                res = f"Your command {cmd} not in allowed: {set(commands)}"
            else:
                debug(f"command: {cmd}, args: {args}")
                res = getattr(self, cmd)(args)
        if res:
            print(res)


if __name__ == "__main__":
    TunnelerCLI().invoke(sys.argv[1:])
