"""
lib/net.py — Connection helpers (p4-team style: nc / receive_until_match / send).

Works with pwntools if installed; otherwise falls back to a raw-socket
implementation so your solvers run even on a bare Python.
"""
import os
import re
import socket
import sys
import time

try:
    from pwn import remote, context, p64, u64, ELF, ROP  # noqa
    HAVE_PWN = True
except Exception:
    HAVE_PWN = False


class Conn:
    """Thin wrapper around pwntools.remote, or a raw socket if pwntools is missing."""

    def __init__(self, host, port, timeout=10, use_pwntools=True):
        self.host = host
        self.port = port
        self.timeout = timeout
        self.use_pwntools = use_pwntools and HAVE_PWN
        if self.use_pwntools:
            context.log_level = os.environ.get("CTF_LOG", "info")
            self.s = remote(host, port, timeout=timeout)
        else:
            self.s = socket.create_connection((host, port), timeout=timeout)
            self._buf = b""

    # ---- low level ----
    def recv_raw(self, n=4096):
        if self.use_pwntools:
            return self.s.recv(n)
        data = b""
        try:
            while len(data) < n:
                chunk = self.s.recv(n - len(data))
                if not chunk:
                    break
                data += chunk
        except socket.timeout:
            pass
        return data

    def recv_until(self, marker, timeout=None):
        """Receive until `marker` (bytes) appears. Returns everything including marker."""
        if self.use_pwntools:
            return self.s.recvuntil(marker)
        marker = marker.encode() if isinstance(marker, str) else marker
        end = time.time() + (timeout or self.timeout)
        buf = self._buf
        while marker not in buf and time.time() < end:
            try:
                self.s.settimeout(max(0.1, end - time.time()))
                chunk = self.s.recv(4096)
                if not chunk:
                    break
                buf += chunk
            except socket.timeout:
                break
        self._buf = b""
        return buf

    def recv_until_match(self, pat, timeout=None):
        """Receive until regex `pat` matches. Returns the matched prefix+match."""
        if self.use_pwntools:
            return self.s.recvline_regex(pat) if hasattr(self.s, "recvline_regex") else self.s.recvuntil(pat.encode())
        rx = re.compile(pat.encode() if isinstance(pat, str) else pat)
        end = time.time() + (timeout or self.timeout)
        buf = self._buf
        while time.time() < end:
            m = rx.search(buf)
            if m:
                self._buf = buf[m.end():]
                return buf[:m.end()]
            try:
                self.s.settimeout(max(0.1, end - time.time()))
                chunk = self.s.recv(4096)
                if not chunk:
                    break
                buf += chunk
            except socket.timeout:
                break
        return buf

    def send(self, data):
        data = data.encode() if isinstance(data, str) else data
        if self.use_pwntools:
            self.s.sendline(data) if data.endswith(b"\n") else self.s.send(data)
        else:
            self.s.sendall(data)

    def sendline(self, data):
        data = data.encode() if isinstance(data, str) else data
        if self.use_pwntools:
            self.s.sendline(data)
        else:
            self.s.sendall(data + b"\n")

    def interactive(self):
        if self.use_pwntools:
            self.s.interactive()
        else:
            print("[!] interactive needs pwntools; dropping to manual mode")
            import select
            while True:
                r, _, _ = select.select([self.s, sys.stdin], [], [])
                if self.s in r:
                    print(self.recv_raw(4096).decode(errors="replace"), end="")
                if sys.stdin in r:
                    self.s.sendall(sys.stdin.readline().encode())

    def close(self):
        try:
            self.s.close()
        except Exception:
            pass


def nc(host, port, timeout=10):
    """Drop-in for the p4 pattern `nc(host, port)`."""
    return Conn(host, port, timeout=timeout)


# convenience re-exports for templates
def receive_until(s, marker, timeout=None):
    return s.recv_until(marker, timeout)


def receive_until_match(s, pat, timeout=None):
    return s.recv_until_match(pat, timeout)


def send(s, data):
    s.send(data)