package chat.vojo.app;

import android.util.Log;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicBoolean;

import chat.vojo.app.ProxyConfigStore.ProxyCreds;

/**
 * Authenticated HTTP-CONNECT relay on 127.0.0.1:&lt;ephemeral&gt; that exists
 * ONLY while the proxy is enabled (proxy_support.md §15 #1/#2/#3).
 *
 * Why it exists: Chromium's WebView proxy can't present SOCKS5 user/pass auth,
 * and the five native HttpURLConnection paths bypass the WebView proxy
 * entirely. Both are pointed at this loopback relay; the relay holds the
 * upstream credentials and dials the user's real proxy (HTTP-CONNECT or SOCKS5,
 * with or without auth).
 *
 * Inbound is authenticated with a high-entropy per-session token via standard
 * proxy Basic auth (407 challenge) — answered by the WebView's
 * onReceivedHttpAuthRequest and, for native paths, a loopback-scoped
 * Authenticator. Without the token any co-resident app could use the loopback
 * listener as an open relay (the §15 #1 HIGH threat). Peer-UID filtering is
 * NOT relied on (unreliable on modern Android).
 *
 * Pure L4 forwarder after CONNECT: no TLS inspection, no CA — the WebView's
 * TLS to the homeserver is end-to-end through the tunnel.
 */
final class LocalProxyRelay {

    private static final String TAG = "LocalProxyRelay";
    private static final String AUTH_USER = "vojo";
    // The Basic realm sent in our 407. The WebView's onReceivedHttpAuthRequest
    // is matched on THIS (not the host) so the proxy-auth answer works whatever
    // host Chromium reports for the relay's 407 (proxy_android_impl.md §1).
    private static final String AUTH_REALM = "vojo";
    private static final int ACCEPT_BACKLOG = 32;
    private static final int UPSTREAM_CONNECT_TIMEOUT_MS = 15000;
    private static final int HANDSHAKE_TIMEOUT_MS = 15000;
    private static final int MAX_HEADER_BYTES = 16 * 1024;

    private final ProxyCreds creds;
    private final String token;
    private final byte[] expectedAuthBytes;

    private final AtomicBoolean running = new AtomicBoolean(false);
    // Every established client+upstream socket, so disable()/logout can
    // force-close in-flight tunnels (e.g. the /sync long-poll) instead of
    // letting them ride the user's proxy until the far end FINs.
    private final java.util.Set<Socket> live =
        java.util.Collections.synchronizedSet(new java.util.HashSet<>());
    private ServerSocket serverSocket;
    private Thread acceptThread;
    private ExecutorService workers;
    private int port = -1;

    LocalProxyRelay(ProxyCreds creds, String token) {
        this.creds = creds;
        this.token = token;
        this.expectedAuthBytes = ("Basic " + base64(AUTH_USER + ":" + token))
            .getBytes(StandardCharsets.US_ASCII);
    }

    int getPort() {
        return port;
    }

    String getToken() {
        return token;
    }

    static String getAuthUser() {
        return AUTH_USER;
    }

    static String getAuthRealm() {
        return AUTH_REALM;
    }

    boolean isRunning() {
        return running.get();
    }

    /** Bind the loopback listener and start accepting. Idempotent-ish: throws
     *  if already started. Must be called off the main thread. */
    synchronized void start() throws IOException {
        if (running.get()) return;
        serverSocket = new ServerSocket(0, ACCEPT_BACKLOG, InetAddress.getByName("127.0.0.1"));
        port = serverSocket.getLocalPort();
        workers = Executors.newCachedThreadPool();
        running.set(true);
        acceptThread = new Thread(this::acceptLoop, "vojo-proxy-relay");
        acceptThread.setDaemon(true);
        acceptThread.start();
        Log.i(TAG, "relay up on 127.0.0.1:" + port);
    }

    synchronized void stop() {
        if (!running.getAndSet(false)) return;
        try {
            if (serverSocket != null) serverSocket.close();
        } catch (IOException ignored) {
            // closing the listener is best-effort
        }
        if (workers != null) workers.shutdownNow();
        if (acceptThread != null) acceptThread.interrupt();
        // Force-close every in-flight tunnel so no traffic keeps riding the
        // user's upstream after the proxy is turned off / the user logs out.
        Socket[] snapshot;
        synchronized (live) {
            snapshot = live.toArray(new Socket[0]);
        }
        for (Socket s : snapshot) {
            closeQuietly(s);
        }
        live.clear();
        serverSocket = null;
        port = -1;
        Log.i(TAG, "relay stopped");
    }

    private void acceptLoop() {
        while (running.get()) {
            final Socket client;
            // Capture the listener locally so a concurrent stop() that nulls
            // serverSocket can't NPE us between the running check and accept().
            final ServerSocket ss = serverSocket;
            if (ss == null) break;
            try {
                client = ss.accept();
            } catch (IOException e) {
                if (running.get()) Log.w(TAG, "accept failed");
                break;
            }
            try {
                workers.execute(() -> handle(client));
            } catch (RuntimeException rejected) {
                closeQuietly(client);
            }
        }
    }

    private void handle(Socket client) {
        Socket upstream = null;
        try {
            client.setSoTimeout(HANDSHAKE_TIMEOUT_MS);
            track(client);
            InputStream cin = client.getInputStream();
            OutputStream cout = client.getOutputStream();

            String[] headerLines = readHeaders(cin);
            if (headerLines == null) {
                closeQuietly(client);
                return;
            }

            // Authenticate FIRST so an unauthenticated probe gets only a 407 —
            // never a fingerprintable 405/400 that confirms what's listening.
            if (!authOk(headerLines)) {
                // 407 challenge — the client (WebView / native Authenticator)
                // retries with the token.
                cout.write(("HTTP/1.1 407 Proxy Authentication Required\r\n"
                    + "Proxy-Authenticate: Basic realm=\"vojo\"\r\n"
                    + "Content-Length: 0\r\n"
                    + "Connection: close\r\n\r\n").getBytes(StandardCharsets.US_ASCII));
                cout.flush();
                closeQuietly(client);
                return;
            }

            String requestLine = headerLines.length > 0 ? headerLines[0] : "";
            String[] parts = requestLine.split(" ");
            if (parts.length < 2 || !"CONNECT".equalsIgnoreCase(parts[0])) {
                writeStatus(cout, "405 Method Not Allowed");
                closeQuietly(client);
                return;
            }

            String[] hostPort = splitHostPort(parts[1]);
            if (hostPort == null) {
                writeStatus(cout, "400 Bad Request");
                closeQuietly(client);
                return;
            }
            String targetHost = hostPort[0];
            int targetPort;
            try {
                targetPort = Integer.parseInt(hostPort[1]);
            } catch (NumberFormatException nfe) {
                writeStatus(cout, "400 Bad Request");
                closeQuietly(client);
                return;
            }

            upstream = dialUpstream(targetHost, targetPort);
            if (upstream == null) {
                writeStatus(cout, "502 Bad Gateway");
                closeQuietly(client);
                return;
            }
            track(upstream);

            cout.write("HTTP/1.1 200 Connection Established\r\n\r\n"
                .getBytes(StandardCharsets.US_ASCII));
            cout.flush();

            // Tunnel established — pure byte pipe. Drop the handshake timeout so
            // a long-poll /sync isn't torn down mid-stream.
            client.setSoTimeout(0);
            upstream.setSoTimeout(0);
            pipe(client, upstream);
        } catch (IOException e) {
            closeQuietly(client);
            closeQuietly(upstream);
        }
    }

    // ── Upstream dialers ────────────────────────────────────────────────────

    private Socket dialUpstream(String targetHost, int targetPort) {
        try {
            if ("http".equalsIgnoreCase(creds.type)) {
                return dialHttpUpstream(targetHost, targetPort);
            }
            return dialSocks5Upstream(targetHost, targetPort);
        } catch (IOException e) {
            Log.w(TAG, "upstream dial failed");
            return null;
        }
    }

    private Socket dialHttpUpstream(String targetHost, int targetPort) throws IOException {
        Socket s = new Socket();
        s.connect(new InetSocketAddress(creds.host, creds.port), UPSTREAM_CONNECT_TIMEOUT_MS);
        s.setSoTimeout(HANDSHAKE_TIMEOUT_MS);
        StringBuilder req = new StringBuilder();
        req.append("CONNECT ").append(targetHost).append(':').append(targetPort)
            .append(" HTTP/1.1\r\n");
        req.append("Host: ").append(targetHost).append(':').append(targetPort).append("\r\n");
        if (creds.hasAuth()) {
            req.append("Proxy-Authorization: Basic ")
                .append(base64(creds.username + ":" + (creds.password == null ? "" : creds.password)))
                .append("\r\n");
        }
        req.append("\r\n");
        s.getOutputStream().write(req.toString().getBytes(StandardCharsets.US_ASCII));
        s.getOutputStream().flush();

        String[] respLines = readHeaders(s.getInputStream());
        if (respLines == null || respLines.length == 0) {
            closeQuietly(s);
            return null;
        }
        // Expect "HTTP/1.x 200 ..." — parse the numeric status token, not a
        // substring (a reason phrase could otherwise contain "200").
        String[] statusParts = respLines[0].split(" ");
        if (statusParts.length < 2 || !"200".equals(statusParts[1])) {
            closeQuietly(s);
            return null;
        }
        return s;
    }

    private Socket dialSocks5Upstream(String targetHost, int targetPort) throws IOException {
        Socket s = new Socket();
        s.connect(new InetSocketAddress(creds.host, creds.port), UPSTREAM_CONNECT_TIMEOUT_MS);
        s.setSoTimeout(HANDSHAKE_TIMEOUT_MS);
        InputStream in = s.getInputStream();
        OutputStream out = s.getOutputStream();

        boolean auth = creds.hasAuth();
        // Greeting: when we have creds advertise ONLY user/pass (0x02) so a
        // hostile/misconfigured upstream can't downgrade us to no-auth and have
        // the creds silently dropped. Otherwise advertise only no-auth.
        if (auth) {
            out.write(new byte[] {0x05, 0x01, 0x02});
        } else {
            out.write(new byte[] {0x05, 0x01, 0x00});
        }
        out.flush();
        byte[] methodSel = readN(in, 2);
        if (methodSel == null || methodSel[0] != 0x05) {
            closeQuietly(s);
            return null;
        }
        int method = methodSel[1] & 0xff;
        if (auth) {
            // We advertised ONLY user/pass — reject any downgrade.
            if (method != 0x02) { closeQuietly(s); return null; }
            // RFC 1929 user/pass sub-negotiation.
            byte[] u = creds.username.getBytes(StandardCharsets.UTF_8);
            byte[] p = (creds.password == null ? "" : creds.password).getBytes(StandardCharsets.UTF_8);
            if (u.length > 255 || p.length > 255) { closeQuietly(s); return null; }
            ByteArrayOutputStream a = new ByteArrayOutputStream();
            a.write(0x01);
            a.write(u.length);
            a.write(u, 0, u.length);
            a.write(p.length);
            a.write(p, 0, p.length);
            out.write(a.toByteArray());
            out.flush();
            byte[] authResp = readN(in, 2);
            if (authResp == null || authResp[1] != 0x00) {
                closeQuietly(s);
                return null;
            }
        } else if (method != 0x00) {
            // We advertised only no-auth; anything else is unacceptable.
            closeQuietly(s);
            return null;
        }

        // CONNECT request with the target as a DOMAINNAME so the upstream does
        // remote DNS (no on-device DNS leak).
        byte[] host = targetHost.getBytes(StandardCharsets.UTF_8);
        if (host.length > 255) { closeQuietly(s); return null; }
        ByteArrayOutputStream req = new ByteArrayOutputStream();
        req.write(0x05); // version
        req.write(0x01); // CONNECT
        req.write(0x00); // reserved
        req.write(0x03); // ATYP = domain
        req.write(host.length);
        req.write(host, 0, host.length);
        req.write((targetPort >> 8) & 0xff);
        req.write(targetPort & 0xff);
        out.write(req.toByteArray());
        out.flush();

        // Reply: VER REP RSV ATYP BND.ADDR BND.PORT — read the fixed prefix
        // then the variable address.
        byte[] head = readN(in, 4);
        if (head == null || head[0] != 0x05 || head[1] != 0x00) {
            closeQuietly(s);
            return null;
        }
        int atyp = head[3] & 0xff;
        int addrLen;
        if (atyp == 0x01) {
            addrLen = 4;
        } else if (atyp == 0x04) {
            addrLen = 16;
        } else if (atyp == 0x03) {
            byte[] l = readN(in, 1);
            if (l == null) { closeQuietly(s); return null; }
            addrLen = l[0] & 0xff;
        } else {
            closeQuietly(s);
            return null;
        }
        if (readN(in, addrLen + 2) == null) { // address + 2-byte port
            closeQuietly(s);
            return null;
        }
        return s;
    }

    // ── Helpers ─────────────────────────────────────────────────────────────

    private boolean authOk(String[] headerLines) {
        for (String line : headerLines) {
            int colon = line.indexOf(':');
            if (colon <= 0) continue;
            String name = line.substring(0, colon).trim();
            if (!"Proxy-Authorization".equalsIgnoreCase(name)) continue;
            String value = line.substring(colon + 1).trim();
            byte[] got = value.getBytes(StandardCharsets.US_ASCII);
            // Constant-time compare to avoid leaking the token via timing.
            return MessageDigest.isEqual(got, expectedAuthBytes);
        }
        return false;
    }

    // Read CRLF-delimited header lines until the blank line. Returns the lines
    // (request/status line first), or null on malformed / oversized input. Reads
    // byte-by-byte so it never consumes tunnel body past the header terminator.
    private static String[] readHeaders(InputStream in) throws IOException {
        ByteArrayOutputStream buf = new ByteArrayOutputStream();
        int prev = -1;
        int total = 0;
        while (true) {
            int b = in.read();
            if (b == -1) {
                if (buf.size() == 0) return null;
                break;
            }
            total++;
            if (total > MAX_HEADER_BYTES) return null;
            buf.write(b);
            if (prev == '\r' && b == '\n') {
                int size = buf.size();
                byte[] arr = buf.toByteArray();
                // Header section ends with CRLFCRLF.
                if (size >= 4 && arr[size - 3] == '\n' && arr[size - 4] == '\r') {
                    break;
                }
            }
            prev = b;
        }
        String text = new String(buf.toByteArray(), StandardCharsets.US_ASCII);
        String[] rawLines = text.split("\r\n");
        return rawLines;
    }

    private static String[] splitHostPort(String authority) {
        int idx = authority.lastIndexOf(':');
        if (idx <= 0 || idx == authority.length() - 1) return null;
        String host = authority.substring(0, idx);
        // Strip IPv6 brackets if present.
        if (host.startsWith("[") && host.endsWith("]")) {
            host = host.substring(1, host.length() - 1);
        }
        return new String[] {host, authority.substring(idx + 1)};
    }

    private static byte[] readN(InputStream in, int n) throws IOException {
        byte[] out = new byte[n];
        int off = 0;
        while (off < n) {
            int r = in.read(out, off, n - off);
            if (r == -1) return null;
            off += r;
        }
        return out;
    }

    private static void writeStatus(OutputStream out, String status) {
        try {
            out.write(("HTTP/1.1 " + status + "\r\nContent-Length: 0\r\n"
                + "Connection: close\r\n\r\n").getBytes(StandardCharsets.US_ASCII));
            out.flush();
        } catch (IOException ignored) {
            // best-effort error reply
        }
    }

    private void pipe(Socket a, Socket b) {
        Thread t1 = new Thread(() -> copy(a, b), "vojo-relay-a2b");
        Thread t2 = new Thread(() -> copy(b, a), "vojo-relay-b2a");
        t1.setDaemon(true);
        t2.setDaemon(true);
        t1.start();
        t2.start();
    }

    private void copy(Socket from, Socket to) {
        byte[] buf = new byte[16 * 1024];
        try {
            InputStream in = from.getInputStream();
            OutputStream out = to.getOutputStream();
            int n;
            while ((n = in.read(buf)) != -1) {
                out.write(buf, 0, n);
                out.flush();
            }
        } catch (IOException ignored) {
            // peer closed / reset — fall through to tear both sides down
        } finally {
            closeQuietly(from);
            closeQuietly(to);
        }
    }

    private static String base64(String s) {
        return android.util.Base64.encodeToString(
            s.getBytes(StandardCharsets.UTF_8), android.util.Base64.NO_WRAP);
    }

    private void track(Socket s) {
        if (s != null) live.add(s);
    }

    private void closeQuietly(Socket s) {
        if (s == null) return;
        live.remove(s);
        try {
            s.close();
        } catch (IOException ignored) {
            // best-effort close
        }
    }
}
