package invalid.lena.scrcpy;

import android.content.Context;

import org.bouncycastle.asn1.ASN1EncodableVector;
import org.bouncycastle.asn1.ASN1Encoding;
import org.bouncycastle.asn1.ASN1Integer;
import org.bouncycastle.asn1.DERBitString;
import org.bouncycastle.asn1.DERSequence;
import org.bouncycastle.asn1.pkcs.PKCSObjectIdentifiers;
import org.bouncycastle.asn1.x500.X500Name;
import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo;
import org.bouncycastle.asn1.x509.TBSCertificate;
import org.bouncycastle.asn1.x509.Time;
import org.bouncycastle.asn1.x509.V3TBSCertificateGenerator;

import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.math.BigInteger;
import java.nio.file.Files;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.PrivateKey;
import java.security.SecureRandom;
import java.security.Signature;
import java.security.cert.Certificate;
import java.security.cert.CertificateFactory;
import java.security.spec.PKCS8EncodedKeySpec;
import java.util.Date;
import java.util.concurrent.TimeUnit;

import io.github.muntashirakon.adb.AbsAdbConnectionManager;
import io.github.muntashirakon.adb.AdbStream;
import io.github.muntashirakon.adb.LocalServices;

// Thin facade over libadb-android. One instance per process: pairs once,
// connects per session, then exposes typed openers for the three streams
// the rest of the app cares about (shell, sync, localabstract).
//
// The keypair lives in filesDir/{adbkey,adbcert}: PKCS#8 DER for the key,
// X.509 DER for the cert. Software-backed RSA - AndroidKeyStore is unusable
// here because libadb-android's auth path uses raw RSA/ECB/NoPadding which
// hardware-backed keys refuse.
public final class Adb extends AbsAdbConnectionManager {

    private static final String DEVICE_NAME = "scrcpy-android";
    private static final String KEY_FILE = "adbkey";
    private static final String CERT_FILE = "adbcert";

    // Per-operation timeout for adb protocol calls (pair/connect/openStream).
    // libadb-android defaults to Long.MAX_VALUE; a hung target would block
    // forever otherwise.
    private static final long OP_TIMEOUT_MS = 15_000L;

    // libadb-android's mApi is documented as "the target's API for protocol
    // negotiation". We don't know the target's API at construction time;
    // pinning to the current Android max means the protocol stays at its
    // latest version, which is what wireless-debugging Android-11+ targets
    // expect. Bump when newer protocol versions ship.
    private static final int  TARGET_API_HINT = android.os.Build.VERSION_CODES.UPSIDE_DOWN_CAKE;

    // Singleton: Main and Mirror both want an Adb; building two would
    // race on the on-disk keypair and waste a RSA generation. Lazily
    // initialised on the FIRST caller, which should be off the UI thread.
    private static volatile Adb instance;

    public static Adb getInstance(Context ctx) throws Exception {
        Adb local = instance;
        if (local != null) return local;
        synchronized (Adb.class) {
            if (instance == null) instance = new Adb(ctx.getApplicationContext());
            return instance;
        }
    }

    private final PrivateKey privateKey;
    private final Certificate certificate;

    private Adb(Context ctx) throws Exception {
        setApi(TARGET_API_HINT);
        setTimeout(OP_TIMEOUT_MS, TimeUnit.MILLISECONDS);
        File keyFile = new File(ctx.getFilesDir(), KEY_FILE);
        File certFile = new File(ctx.getFilesDir(), CERT_FILE);
        if (keyFile.exists() && certFile.exists()) {
            privateKey = loadKey(keyFile);
            certificate = loadCert(certFile);
            Log.i("adb: loaded keypair from %s", keyFile);
        } else {
            KeyPair kp = generateRsa();
            privateKey = kp.getPrivate();
            certificate = selfSignedCert(kp);
            saveKey(keyFile, privateKey);
            saveCert(certFile, certificate);
            Log.i("adb: generated new keypair at %s", keyFile);
        }
    }

    @Override protected PrivateKey getPrivateKey()  { return privateKey; }
    @Override protected Certificate getCertificate() { return certificate; }
    @Override protected String      getDeviceName() { return DEVICE_NAME; }

    // ---- typed openers ----

    public AdbStream openAbstract(String name) throws IOException, InterruptedException {
        return openStream(LocalServices.LOCAL_UNIX_SOCKET_ABSTRACT, name);
    }

    public AdbStream openShell(String cmd) throws IOException, InterruptedException {
        // Bypass LocalServices.getDestination(SHELL, args) - it wraps any
        // arg containing a space in literal double quotes, and adbd then
        // tries to exec `"the whole quoted thing"` as a single filename.
        // Build the destination ourselves so the cmd reaches sh -c
        // unmolested.
        return openStream("shell:" + cmd);
    }

    public AdbStream openSync() throws IOException, InterruptedException {
        return openStream(LocalServices.SYNC);
    }

    // ---- key + cert I/O ----

    private static KeyPair generateRsa() throws Exception {
        KeyPairGenerator gen = KeyPairGenerator.getInstance("RSA");
        gen.initialize(2048, new SecureRandom());
        return gen.generateKeyPair();
    }

    private static PrivateKey loadKey(File f) throws Exception {
        byte[] data = Files.readAllBytes(f.toPath());
        return KeyFactory.getInstance("RSA").generatePrivate(new PKCS8EncodedKeySpec(data));
    }

    private static Certificate loadCert(File f) throws Exception {
        try (InputStream in = new FileInputStream(f)) {
            return CertificateFactory.getInstance("X.509").generateCertificate(in);
        }
    }

    private static void saveKey(File f, PrivateKey k) throws IOException {
        AtomicFiles.write(f, k.getEncoded());
    }

    private static void saveCert(File f, Certificate c) throws Exception {
        AtomicFiles.write(f, c.getEncoded());
    }

    // Build a minimal self-signed X.509 certificate over the keypair.
    //
    // ASN.1 layout (RFC 5280 sec. 4.1):
    //
    //   Certificate ::= SEQUENCE {
    //     tbsCertificate       TBSCertificate,
    //     signatureAlgorithm   AlgorithmIdentifier,    -- sha256WithRSAEnc
    //     signature            BIT STRING               -- RSA over tbs
    //   }
    //
    //   TBSCertificate ::= SEQUENCE {
    //     version              [0] EXPLICIT v3,
    //     serialNumber         INTEGER 1,
    //     signatureAlgorithm   AlgorithmIdentifier,
    //     issuer  = subject    Name (CN=scrcpy-android),
    //     validity             { notBefore, notAfter },
    //     subjectPublicKeyInfo SubjectPublicKeyInfo
    //   }
    //
    // ADB validates the peer by the raw RSA public-key fingerprint in the
    // SPKI bits, not by the DN or signature, so the surrounding cert is
    // cosmetic - but a valid X.509 wrapper is still required for the TLS
    // handshake that wireless-debugging uses post-pairing.
    private static Certificate selfSignedCert(KeyPair kp) throws Exception {
        long now = System.currentTimeMillis();
        Date notBefore = new Date(now - 60_000L);
        Date notAfter  = new Date(now + 50L * 365 * 24 * 3600 * 1000L);
        X500Name dn = new X500Name("CN=" + DEVICE_NAME);
        AlgorithmIdentifier sigAlg =
                new AlgorithmIdentifier(PKCSObjectIdentifiers.sha256WithRSAEncryption);

        V3TBSCertificateGenerator g = new V3TBSCertificateGenerator();
        g.setSerialNumber(new ASN1Integer(BigInteger.ONE));
        g.setSignature(sigAlg);
        g.setIssuer(dn);
        g.setSubject(dn);
        g.setStartDate(new Time(notBefore));
        g.setEndDate(new Time(notAfter));
        g.setSubjectPublicKeyInfo(
                SubjectPublicKeyInfo.getInstance(kp.getPublic().getEncoded()));

        TBSCertificate tbs = g.generateTBSCertificate();

        Signature s = Signature.getInstance("SHA256withRSA");
        s.initSign(kp.getPrivate());
        s.update(tbs.getEncoded(ASN1Encoding.DER));
        byte[] sig = s.sign();

        ASN1EncodableVector v = new ASN1EncodableVector();
        v.add(tbs);
        v.add(sigAlg);
        v.add(new DERBitString(sig));
        byte[] certDer = new DERSequence(v).getEncoded(ASN1Encoding.DER);

        return CertificateFactory.getInstance("X.509")
                .generateCertificate(new ByteArrayInputStream(certDer));
    }
}