package com.elphel.imagej.cuas.rt;

import java.awt.Rectangle;
import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.File;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.nio.ByteBuffer;
import java.nio.file.Files;

/**
 * TCP client for the DGX remote DNN inference server (attic/imagej-elphel-internal/c5p_dnn/infer_server.py).
 * By Claude on 06/20/2026.
 *
 * Stateful, big-endian protocol (Java DataInput/OutputStream is big-endian; the Python side uses
 * struct '>' / numpy '>f4', so they match without byte-swapping):
 *   UPLOAD  : cmd=1, int T,H,W, then T*H*W float32 (the SUBAVG+LoG-conditioned stack)
 *             -> reply: int n_levels, n_levels x int frames_per_level, int N, double build_ms
 *             The DGX builds the temporal pyramid (0.5*(now+prev), == temporalAverageLReLU).
 *   INFER   : cmd=2, int level, newest, roi_x, roi_y, roi_w, roi_h
 *             -> reply: double gpu_ms, int H,W, 5*H*W float32 (offset5 dx,dy,s,Vx,Vy plane-major),
 *                       int rh,rw,nvel, rh*rw*nvel float32 (ROI softmax*s, pixel-major)
 *   READBACK: cmd=3, int level, frame -> reply: int H,W, H*W float32 (one conditioned pyramid frame)
 *   BYE     : cmd=0
 */
public class CuasDnnRemote implements AutoCloseable {
	private static final int CMD_BYE = 0, CMD_UPLOAD = 1, CMD_INFER = 2, CMD_READBACK = 3;
	private final Socket sock;
	private final DataInputStream  in;
	private final DataOutputStream out;
	private int nFrames = 0;   // temporal depth N, learned from the UPLOAD reply

	public CuasDnnRemote(String hostPort) throws Exception {
		String [] hp = hostPort.split(":");
		int port = (hp.length > 1) ? Integer.parseInt(hp[1].trim()) : 5577;
		sock = new Socket(hp[0].trim(), port);
		sock.setTcpNoDelay(true);
		in  = new DataInputStream (new BufferedInputStream (sock.getInputStream(),  1 << 20));
		out = new DataOutputStream(new BufferedOutputStream(sock.getOutputStream(), 1 << 20));
	}

	public int getNFrames() { return nFrames; }

	/** Upload the conditioned stack frames[T][H*W] (row-major) once; the DGX builds the pyramid.
	 *  Returns the per-level frame counts (and sets getNFrames()). */
	public int [] upload(float [][] frames, int H, int W) throws Exception {
		int T = frames.length;
		out.writeInt(CMD_UPLOAD); out.writeInt(T); out.writeInt(H); out.writeInt(W);
		byte [] buf = new byte [H * W * 4];                 // one frame, reused
		ByteBuffer bb = ByteBuffer.wrap(buf);               // big-endian by default
		for (int t = 0; t < T; t++) {
			bb.clear();
			float [] fr = frames[t];
			for (int p = 0; p < fr.length; p++) bb.putFloat(fr[p]);
			out.write(buf);
		}
		out.flush();
		int nl = in.readInt();
		int [] counts = new int [nl];
		for (int i = 0; i < nl; i++) counts[i] = in.readInt();
		nFrames = in.readInt();
		in.readDouble();                                    // build_ms (logged server-side)
		return counts;
	}

	/** One INFER result: full-frame offset5 {dx,dy,s,Vx,Vy}[5][H*W] (plane-major) +
	 *  ROI field [rh*rw][nvel] (softmax*s, pixel-major) + the GPU compute ms. */
	public static class Result {
		public double gpuMs;
		public int H, W, rh, rw, nvel;
		public float [][] offset5;    // [5][H*W]
		public float [][] roiField;   // [rh*rw][nvel]
	}

	public Result infer(int level, int newest, Rectangle roi) throws Exception {
		out.writeInt(CMD_INFER); out.writeInt(level); out.writeInt(newest);
		out.writeInt(roi.x); out.writeInt(roi.y); out.writeInt(roi.width); out.writeInt(roi.height);
		out.flush();
		Result r = new Result();
		r.gpuMs = in.readDouble(); r.H = in.readInt(); r.W = in.readInt();
		int hw = r.H * r.W;
		byte [] ob = new byte [5 * hw * 4]; in.readFully(ob);
		ByteBuffer obb = ByteBuffer.wrap(ob);
		r.offset5 = new float [5][hw];
		for (int c = 0; c < 5; c++) for (int p = 0; p < hw; p++) r.offset5[c][p] = obb.getFloat();
		r.rh = in.readInt(); r.rw = in.readInt(); r.nvel = in.readInt();
		int rn = r.rh * r.rw;
		byte [] rb = new byte [rn * r.nvel * 4]; in.readFully(rb);
		ByteBuffer rbb = ByteBuffer.wrap(rb);
		r.roiField = new float [rn][r.nvel];
		for (int p = 0; p < rn; p++) for (int v = 0; v < r.nvel; v++) r.roiField[p][v] = rbb.getFloat();
		return r;
	}

	/** Debug: read back one conditioned pyramid frame [H*W] (sets hw[0]=H, hw[1]=W if provided). */
	public float [] readback(int level, int frame, int [] hw) throws Exception {
		out.writeInt(CMD_READBACK); out.writeInt(level); out.writeInt(frame); out.flush();
		int H = in.readInt(), W = in.readInt();
		if ((hw != null) && (hw.length >= 2)) { hw[0] = H; hw[1] = W; }
		byte [] b = new byte [H * W * 4]; in.readFully(b);
		ByteBuffer bb = ByteBuffer.wrap(b);
		float [] f = new float [H * W];
		for (int i = 0; i < f.length; i++) f[i] = bb.getFloat();
		return f;
	}

	@Override
	public void close() {
		try { out.writeInt(CMD_BYE); out.flush(); } catch (Exception e) { /* ignore */ }
		try { sock.close(); } catch (Exception e) { /* ignore */ }
	}

	// ---- Auto-launch (no manual interaction): ensure the DGX server is running, else deploy + start it ----
	// Mirrors the GPU-kernel default-vs-override: server scripts come from the bundled jar resource
	// (/cuas_dnn/...) unless srcdir is set, then from that local dir. By Claude on 06/20/2026

	private static boolean canConnect(String host, int port, int timeoutMs) {
		try (Socket s = new Socket()) { s.connect(new InetSocketAddress(host, port), timeoutMs); return true; }
		catch (Exception e) { return false; }
	}

	/** Read a server script: bundled jar resource (/cuas_dnn/<name>) by default, or <srcdir>/<name> if set. */
	private static byte [] readServerScript(String name, String srcdir) throws Exception {
		if ((srcdir != null) && !srcdir.isEmpty()) {
			return Files.readAllBytes(new File(srcdir, name).toPath());           // local-repo override (config)
		}
		try (InputStream is = CuasDnnRemote.class.getResourceAsStream("/cuas_dnn/" + name)) {
			if (is == null) throw new Exception("bundled resource /cuas_dnn/" + name + " not found");
			return is.readAllBytes();                                             // working version in jar
		}
	}

	private static void deployScript(String name, String srcdir, String sshTarget, String dest) throws Exception {
		byte [] content = readServerScript(name, srcdir);
		ProcessBuilder pb = new ProcessBuilder("ssh", sshTarget, "cat > " + dest);
		pb.redirectErrorStream(true);
		Process p = pb.start();
		try (OutputStream os = p.getOutputStream()) { os.write(content); }
		if (p.waitFor() != 0) throw new Exception("deploy " + name + " to " + sshTarget + ":" + dest + " failed");
	}

	/** Ensure the DGX server is reachable at hostPort; if not, deploy the (bundled/override) scripts and
	 *  ssh-launch run_infer_server.sh with RUN=model, then poll until it accepts connections. No manual step. */
	public static void ensureServer(String hostPort, String model, String srcdir) throws Exception {
		String [] hp = hostPort.split(":");
		String host = hp[0].trim();
		int port = (hp.length > 1) ? Integer.parseInt(hp[1].trim()) : 5577;
		if (canConnect(host, port, 1500)) return;                                 // already up
		String sshTarget = "elphel@" + host;                                      // DGX login user
		String code = "/home/elphel/c5p_dnn";                                     // DGX dir (model.py + runs/ live here)
		System.out.println("CuasDnnRemote.ensureServer(): no server at " + host + ":" + port
				+ " - deploying scripts + launching on " + sshTarget + " (model=" + model + ")");
		deployScript("infer_server.py",     srcdir, sshTarget, code + "/infer_server.py");
		deployScript("run_infer_server.sh",  srcdir, sshTarget, code + "/run_infer_server.sh");
		ProcessBuilder pb = new ProcessBuilder("ssh", sshTarget,
				"cd " + code + " && chmod +x run_infer_server.sh && RUN=" + model + " PORT=" + port + " ./run_infer_server.sh start");
		pb.inheritIO();
		pb.start().waitFor();
		long deadline = System.currentTimeMillis() + 90000;                       // model load + warm-up can take a bit
		while (System.currentTimeMillis() < deadline) {
			if (canConnect(host, port, 1500)) { System.out.println("CuasDnnRemote.ensureServer(): server up at " + host + ":" + port); return; }
			Thread.sleep(2000);
		}
		throw new Exception("CuasDnnRemote.ensureServer(): server did not come up at " + host + ":" + port);
	}
}
