"""SOCKS5-Proxy (RFC 1928/1929): CONNECT (TCP) und UDP ASSOCIATE."""
import asyncio
import ipaddress
import socket
import struct
import time

from . import config
from .auth import users
from .core import Blocked, Conn, check_and_resolve, history, log, open_upstream, register, relay
from .policy import policy

REP_OK, REP_FAIL, REP_NOT_ALLOWED, REP_NET_UNREACH, REP_HOST_UNREACH, REP_REFUSED, REP_TTL, REP_CMD, REP_ATYP = range(9)

_udp_in_use: set[int] = set()


def _reply(rep: int, bind_ip: str = "0.0.0.0", bind_port: int = 0) -> bytes:
    ip = ipaddress.ip_address(bind_ip)
    atyp = 1 if ip.version == 4 else 4
    return bytes([5, rep, 0, atyp]) + ip.packed + struct.pack("!H", bind_port)


def _decode_host(raw: bytes) -> str:
    try:
        return raw.decode("idna")
    except UnicodeError:
        return raw.decode("ascii", "replace")


async def _read_addr(reader, atyp: int) -> str:
    if atyp == 1:
        return socket.inet_ntop(socket.AF_INET, await reader.readexactly(4))
    if atyp == 4:
        return socket.inet_ntop(socket.AF_INET6, await reader.readexactly(16))
    if atyp == 3:
        n = (await reader.readexactly(1))[0]
        return _decode_host(await reader.readexactly(n))
    raise ValueError("Adresstyp")


def _rep_for(exc: Exception) -> int:
    if isinstance(exc, ConnectionRefusedError):
        return REP_REFUSED
    if "DNS" in str(exc):
        return REP_HOST_UNREACH
    if "Zeitüberschreitung" in str(exc):
        return REP_TTL
    return REP_HOST_UNREACH


async def handle_socks5(reader: asyncio.StreamReader, writer: asyncio.StreamWriter, client_ip: str) -> None:
    # Begrüßung (Versionsbyte bereits gelesen)
    n = (await reader.readexactly(1))[0]
    methods = await reader.readexactly(n)
    user = None
    if 2 in methods:
        writer.write(b"\x05\x02")
        await writer.drain()
        ver = (await reader.readexactly(1))[0]
        ulen = (await reader.readexactly(1))[0]
        username = (await reader.readexactly(ulen)).decode("utf-8", "replace")
        plen = (await reader.readexactly(1))[0]
        password = (await reader.readexactly(plen)).decode("utf-8", "replace")
        user = await users.authenticate(username, password, client_ip) if ver == 1 else None
        if not user:
            history.add(username=username[:100], client_ip=client_ip, proto="socks", status="auth_failed",
                        detail="Anmeldung fehlgeschlagen")
            writer.write(b"\x01\x01")
            await writer.drain()
            return
        writer.write(b"\x01\x00")
    elif 0 in methods and policy.is_anonymous_allowed(client_ip):
        writer.write(b"\x05\x00")
    else:
        writer.write(b"\x05\xff")
        await writer.drain()
        return
    await writer.drain()
    username = user["username"] if user else "anonym"

    ver, cmd, _, atyp = await reader.readexactly(4)
    try:
        host = await _read_addr(reader, atyp)
    except ValueError:
        writer.write(_reply(REP_ATYP))
        return
    port = struct.unpack("!H", await reader.readexactly(2))[0]

    if cmd == 1:
        await _connect(reader, writer, user, username, client_ip, host, port)
    elif cmd == 3:
        await _udp_associate(reader, writer, user, username, client_ip)
    else:
        writer.write(_reply(REP_CMD))
        await writer.drain()


async def _connect(reader, writer, user, username, client_ip, host, port) -> None:
    conn = Conn(username, client_ip, "socks", host, port, "CONNECT", f"{host}:{port}")
    try:
        r_reader, r_writer = await open_upstream(user, host, port)
    except Blocked as b:
        conn.finish("blocked", b.decision.reason)
        writer.write(_reply(REP_NOT_ALLOWED))
        await writer.drain()
        return
    except OSError as e:
        conn.finish("error", str(e))
        writer.write(_reply(_rep_for(e)))
        await writer.drain()
        return
    sock = r_writer.get_extra_info("sockname")
    try:
        writer.write(_reply(REP_OK, sock[0].split("%")[0], sock[1]))
    except ValueError:
        writer.write(_reply(REP_OK))
    register(conn)
    try:
        await writer.drain()
        await relay(conn, reader, writer, r_reader, r_writer)
    finally:
        conn.finish("ok")


# --------------------------------------------------------------------------- UDP
class UdpRelay(asyncio.DatagramProtocol):
    def __init__(self, user, username, client_ip):
        self.user = user
        self.username = username
        self.client_ip = client_ip
        self.client_addr = None
        self.transport = None
        self.flows: dict[tuple, dict] = {}       # (host, port) -> Statistik
        self.addr_map: dict[tuple, tuple] = {}   # (ip, port) -> (host, port) für Antworten
        self.decisions: dict[tuple, object] = {}
        self.pending: set = set()

    def connection_made(self, transport):
        self.transport = transport

    def datagram_received(self, data: bytes, addr):
        ip = addr[0][7:] if addr[0].startswith("::ffff:") else addr[0]
        if ip == self.client_ip and (self.client_addr is None or addr == self.client_addr) \
                and (addr[0], addr[1]) not in self.addr_map:
            self.client_addr = addr
            t = asyncio.ensure_future(self._from_client(data))
            self.pending.add(t)
            t.add_done_callback(self.pending.discard)
        else:
            self._from_remote(data, addr)

    async def _from_client(self, data: bytes):
        try:
            if len(data) < 4 or data[2] != 0:   # Fragmentierung nicht unterstützt
                return
            atyp = data[3]
            pos = 4
            if atyp == 1:
                host = socket.inet_ntop(socket.AF_INET, data[pos:pos + 4]); pos += 4
            elif atyp == 4:
                host = socket.inet_ntop(socket.AF_INET6, data[pos:pos + 16]); pos += 16
            elif atyp == 3:
                n = data[pos]; host = _decode_host(data[pos + 1:pos + 1 + n]); pos += 1 + n
            else:
                return
            port = struct.unpack("!H", data[pos:pos + 2])[0]
            payload = data[pos + 2:]
            key = (host, port)
            flow = self.flows.get(key)
            if flow is None:
                flow = self.flows[key] = {"start": time.time(), "up": 0, "down": 0, "status": "ok", "detail": "",
                                          "target": None}
                try:
                    addrs = await check_and_resolve(self.user, host, port, socket.SOCK_DGRAM)
                    family = self.transport.get_extra_info("socket").family
                    pick = next((a for f, a in addrs if f == family or family == socket.AF_INET6), None)
                    if pick is None:
                        raise OSError("Keine passende Adressfamilie")
                    target = (pick[0], port)
                    if family == socket.AF_INET6 and ":" not in pick[0]:
                        target = ("::ffff:" + pick[0], port)
                    flow["target"] = target
                    self.addr_map[target] = key
                except Blocked as b:
                    flow["status"], flow["detail"] = "blocked", b.decision.reason
                except OSError as e:
                    flow["status"], flow["detail"] = "error", str(e)
            if flow["target"] is None:
                return
            self.transport.sendto(payload, flow["target"])
            flow["up"] += len(payload)
        except Exception as e:  # noqa: BLE001
            log.debug("UDP-Relay: %s", e)

    def _from_remote(self, data: bytes, addr):
        key = self.addr_map.get((addr[0], addr[1]))
        if key is None or self.client_addr is None:
            return
        self.flows[key]["down"] += len(data)
        host, port = addr[0], addr[1]
        if host.startswith("::ffff:"):
            host = host[7:]
        ip = ipaddress.ip_address(host)
        header = b"\x00\x00\x00" + (b"\x01" if ip.version == 4 else b"\x04") + ip.packed + struct.pack("!H", port)
        self.transport.sendto(header + data, self.client_addr)

    def finish(self):
        for (host, port), f in self.flows.items():
            history.add(ts=f["start"], username=self.username, client_ip=self.client_ip, proto="socks-udp",
                        method="UDP", host=host, port=port, url=f"{host}:{port}", status=f["status"],
                        detail=f["detail"], bytes_up=f["up"], bytes_down=f["down"],
                        duration=round(time.time() - f["start"], 3))


async def _udp_associate(reader, writer, user, username, client_ip) -> None:
    loop = asyncio.get_running_loop()
    transport = proto = None
    bound_port = 0
    for p in config.SOCKS_UDP_PORTS:
        if p in _udp_in_use:
            continue
        try:
            transport, proto = await loop.create_datagram_endpoint(
                lambda: UdpRelay(user, username, client_ip), local_addr=("::", p), family=socket.AF_INET6)
        except OSError:
            try:
                transport, proto = await loop.create_datagram_endpoint(
                    lambda: UdpRelay(user, username, client_ip), local_addr=("0.0.0.0", p))
            except OSError:
                continue
        bound_port = p
        _udp_in_use.add(p)
        break
    if transport is None:
        writer.write(_reply(REP_FAIL))
        await writer.drain()
        return

    advertise = config.SOCKS_UDP_ADVERTISE
    if not advertise:
        local = writer.get_extra_info("sockname")[0]
        advertise = local[7:] if local.startswith("::ffff:") else local.split("%")[0]
    try:
        ipaddress.ip_address(advertise)
    except ValueError:
        try:
            advertise = socket.gethostbyname(advertise)
        except OSError:
            advertise = "0.0.0.0"

    conn = register(Conn(username, client_ip, "socks-udp", "*", bound_port, "UDP ASSOCIATE",
                         f"UDP-Relay Port {bound_port}"))
    conn.closers.append(writer.close)
    try:
        writer.write(_reply(REP_OK, advertise, bound_port))
        await writer.drain()
        # Assoziation lebt, solange die TCP-Verbindung offen ist
        while True:
            data = await reader.read(1024)
            if not data:
                break
            conn.bytes_up = sum(f["up"] for f in proto.flows.values())
            conn.bytes_down = sum(f["down"] for f in proto.flows.values())
    except (ConnectionError, OSError):
        pass
    finally:
        transport.close()
        _udp_in_use.discard(bound_port)
        proto.finish()
        from .core import registry
        registry.pop(conn.id, None)