package com.foxx.androidcast.remoteaccess.control;

import android.util.Log;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.Socket;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.Base64;
import java.util.concurrent.atomic.AtomicBoolean;

/** Minimal WebSocket text channel driving {@code /system/bin/sh} in the app sandbox. */
final class RemoteAccessControlShellHandler implements Runnable {
    private static final String TAG = "RaControlShell";
    private static final String WS_GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";

    private final Socket client;
    private final String requestHeaders;

    RemoteAccessControlShellHandler(Socket client, String requestHeaders) {
        this.client = client;
        this.requestHeaders = requestHeaders != null ? requestHeaders : "";
    }

    @Override
    public void run() {
        try {
            if (!performHandshake()) {
                return;
            }
            runShell();
        } catch (Exception e) {
            Log.d(TAG, "shell session ended: " + e.getMessage());
        } finally {
            try {
                client.close();
            } catch (IOException ignored) {
            }
        }
    }

    private boolean performHandshake() throws IOException {
        String key = headerValue("Sec-WebSocket-Key");
        if (key.isEmpty()) {
            writeRaw("HTTP/1.1 400 Bad Request\r\nConnection: close\r\n\r\n");
            return false;
        }
        String accept = computeAcceptKey(key);
        String response = "HTTP/1.1 101 Switching Protocols\r\n"
                + "Upgrade: websocket\r\n"
                + "Connection: Upgrade\r\n"
                + "Sec-WebSocket-Accept: " + accept + "\r\n\r\n";
        writeRaw(response);
        return true;
    }

    private void runShell() throws IOException, InterruptedException {
        Process process = new ProcessBuilder("/system/bin/sh")
                .redirectErrorStream(true)
                .start();
        AtomicBoolean open = new AtomicBoolean(true);
        Thread reader = new Thread(() -> pumpProcessToClient(process.getInputStream(), open), "RaShellOut");
        reader.start();
        InputStream in = client.getInputStream();
        OutputStream processIn = process.getOutputStream();
        while (open.get()) {
            String text = readTextFrame(in);
            if (text == null) {
                break;
            }
            if ("\u0004".equals(text)) {
                break;
            }
            processIn.write(text.getBytes(StandardCharsets.UTF_8));
            processIn.flush();
        }
        open.set(false);
        process.destroy();
        reader.join(2_000L);
        sendCloseFrame();
    }

    private void pumpProcessToClient(InputStream processOut, AtomicBoolean open) {
        byte[] buf = new byte[4096];
        try {
            int read;
            while (open.get() && (read = processOut.read(buf)) >= 0) {
                if (read == 0) {
                    continue;
                }
                String chunk = new String(buf, 0, read, StandardCharsets.UTF_8);
                sendTextFrame(chunk);
            }
        } catch (IOException e) {
            open.set(false);
        }
    }

    private String headerValue(String name) {
        for (String line : requestHeaders.split("\r\n")) {
            int colon = line.indexOf(':');
            if (colon <= 0) {
                continue;
            }
            if (name.equalsIgnoreCase(line.substring(0, colon).trim())) {
                return line.substring(colon + 1).trim();
            }
        }
        return "";
    }

    private void writeRaw(String text) throws IOException {
        client.getOutputStream().write(text.getBytes(StandardCharsets.US_ASCII));
        client.getOutputStream().flush();
    }

    private synchronized void sendTextFrame(String text) throws IOException {
        byte[] payload = text.getBytes(StandardCharsets.UTF_8);
        OutputStream out = client.getOutputStream();
        out.write(0x81);
        if (payload.length <= 125) {
            out.write(payload.length);
        } else if (payload.length <= 65535) {
            out.write(126);
            out.write((payload.length >> 8) & 0xFF);
            out.write(payload.length & 0xFF);
        } else {
            throw new IOException("frame too large");
        }
        out.write(payload);
        out.flush();
    }

    private synchronized void sendCloseFrame() {
        try {
            OutputStream out = client.getOutputStream();
            out.write(new byte[] {(byte) 0x88, 0x00});
            out.flush();
        } catch (IOException ignored) {
        }
    }

    private static String readTextFrame(InputStream in) throws IOException {
        int b0 = in.read();
        if (b0 < 0) {
            return null;
        }
        int opcode = b0 & 0x0F;
        if (opcode == 0x8) {
            return null;
        }
        int b1 = in.read();
        if (b1 < 0) {
            return null;
        }
        boolean masked = (b1 & 0x80) != 0;
        int len = b1 & 0x7F;
        if (len == 126) {
            len = (in.read() << 8) | in.read();
        } else if (len == 127) {
            for (int i = 0; i < 8; i++) {
                in.read();
            }
            return null;
        }
        byte[] mask = new byte[4];
        if (masked) {
            if (in.read(mask) < 4) {
                return null;
            }
        }
        byte[] payload = new byte[len];
        int off = 0;
        while (off < len) {
            int read = in.read(payload, off, len - off);
            if (read < 0) {
                return null;
            }
            off += read;
        }
        if (masked) {
            for (int i = 0; i < payload.length; i++) {
                payload[i] ^= mask[i % 4];
            }
        }
        return new String(payload, StandardCharsets.UTF_8);
    }

    static String computeAcceptKey(String key) {
        try {
            MessageDigest md = MessageDigest.getInstance("SHA-1");
            md.update((key + WS_GUID).getBytes(StandardCharsets.US_ASCII));
            return Base64.getEncoder().encodeToString(md.digest());
        } catch (NoSuchAlgorithmException e) {
            throw new IllegalStateException(e);
        }
    }
}
